/*
 * ESP8266 Ozdag Attacker V2
 *
 * A redesigned WiFi testing tool for ESP8266.
 * Starting fresh with a new structure.
 */

#include <Arduino.h>
#include <ESP8266WiFi.h>
#include "src/config.h"
#include "src/ui.h"
#include "src/utils.h"
// Include headers for attack modules
#include "src/scanner.h"
#include "src/deauth.h"
#include "src/beacon.h"
#include "src/mitm.h"
#include "src/rogue_ap.h"

void setup() {
    Serial.begin(115200);
    Serial.println("\n\nESP8266 Ozdag Attacker V2 Booting...");

    // Initialize UI (OLED, Button)
    Serial.println("Initializing OLED...");
    if (!oled.begin()) {
        Serial.println("!!! OLED initialization failed!");
        // Optional: loop forever or indicate error differently
        while(1) delay(100);
    } else {
        Serial.println("OLED Initialized OK.");
        
        // --- Start Matrix Boot Animation ---
        const char* bootText = "Ozdag Attacker V2";
        int textWidth = 0;
        int textHeight = 10; // Approx height for 6x10 font
        int textX, textY;

        const int NUM_STREAMS = 20; // Number of falling character streams
        int streamY[NUM_STREAMS];
        int streamX[NUM_STREAMS];
        int streamSpeed[NUM_STREAMS];
        const int streamLen = 5; // How many chars trail behind the leader
        const int charHeight = 7; // Approx height for 5x7 font

        int animDuration = 3000; // Total animation duration in ms
        int stepDelay = 50;      // Delay between frames
        unsigned long animStartTime = millis();

        // Initialize streams
        for (int i = 0; i < NUM_STREAMS; ++i) {
            streamX[i] = random(0, SCREEN_WIDTH / 5) * 5; // Align to 5px grid roughly
            streamY[i] = random(-SCREEN_HEIGHT * 2, 0); // Start off-screen
            streamSpeed[i] = random(2, 6); // Random speeds
        }

        oled.setFont(u8g2_font_6x10_tf); // Font for main text
        textWidth = oled.getStrWidth(bootText);
        textX = (SCREEN_WIDTH - textWidth) / 2;
        textY = SCREEN_HEIGHT - textHeight - 5; // Position text towards bottom

        while (millis() - animStartTime < animDuration) {
            oled.clearBuffer();
            oled.setFont(u8g2_font_5x7_tf); // Font for rain

            // Draw and update streams
            for (int i = 0; i < NUM_STREAMS; ++i) {
                // Draw stream tail (dimmer or different chars could be used)
                for (int j = 1; j <= streamLen; ++j) {
                    int y = streamY[i] - (j * charHeight);
                    if (y > 0 && y < SCREEN_HEIGHT) {
                         // Use a random character for the rain
                        char randomChar = (char)random(33, 127); // Printable ASCII
                        oled.drawGlyph(streamX[i], y, randomChar);
                    }
                }
                
                // Draw leading character (brighter/different? just random for now)
                 if (streamY[i] > 0 && streamY[i] < SCREEN_HEIGHT) {
                    char randomChar = (char)random(33, 127);
                    oled.drawGlyph(streamX[i], streamY[i], randomChar);
                 }

                // Update position
                streamY[i] += streamSpeed[i];

                // Reset if off screen
                if (streamY[i] - (streamLen * charHeight) > SCREEN_HEIGHT) {
                    streamX[i] = random(0, SCREEN_WIDTH / 5) * 5;
                    streamY[i] = random(-charHeight * streamLen, 0); // Reset near top
                    streamSpeed[i] = random(2, 6);
                }
            }

            // --- Draw Static Text --- 
            oled.setFont(u8g2_font_6x10_tf); // Switch font for title
            oled.setDrawColor(1); // Ensure foreground color
            // Optional: Draw a black box behind text for better visibility
            // oled.setDrawColor(0);
            // oled.drawBox(textX - 1, textY - textHeight, textWidth + 2, textHeight + 2);
            // oled.setDrawColor(1);
            oled.drawStr(textX, textY, bootText);

            oled.sendBuffer();
            delay(stepDelay);
        }
        // --- End Matrix Boot Animation ---
        
        // Clear buffer before showing next init messages if desired
        // oled.clearBuffer(); 
        // oled.sendBuffer(); 
    }

    ui_init();
    Serial.println("UI Initialized.");

    // Initialize WiFi in a known state (STA mode, disconnected)
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    Serial.println("WiFi Initialized (STA Mode, Disconnected).");

    // Initialize attack/scanner modules (call their init functions here)
    // Note: Initialization order might matter later depending on dependencies
    scanner_init(); // <<< RESTORE
    deauth_init(); // <<< RESTORE
    beacon_init(); // <<< RESTORE
    // dos_init(); // Removed
    mitm_init(); // <<< UNCOMMENT
    rogue_ap_init(); // <<< UNCOMMENT/RESTORE
    Serial.println("Modules Initialized."); // <<< RESTORE

    Serial.println("Setup Complete.");
}

void loop() {
    // Button handling and drawing are now managed within ui_update()
    ui_update();

    // Update based on current UI state
    switch (currentState) {
        case UIState::MAIN_MENU:
            // Nothing module-specific runs here
            break;
        case UIState::SCANNER_MODE:
            scanner_update(); // <<< RESTORE
            break;
        case UIState::DEAUTH_MODE:
            deauth_update(); // <<< RESTORE
            break;
        case UIState::BEACON_MODE:
            beacon_update(); // <<< RESTORE
            break;
        case UIState::MITM_MODE:
            mitm_update(); // <<< UNCOMMENT
            break;
        case UIState::ROGUE_AP_MODE:
            rogue_ap_update(); // <<< UNCOMMENT/RESTORE
            break;
        default:
            // Handle unknown state? Maybe just yield.
            break;
    }

    yield(); // Allow background tasks
}