#include <stdio.h>
#include <string.h>
#include <inttypes.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"

#include "esp_system.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_wifi.h"

#include "nvs_flash.h"
#include "esp_netif.h"
#include "lwip/inet.h"

#include "led_strip.h"

// Custom Components
#include "iperf.h"
#include "wifi_cfg.h"
#include "csi_log.h"
#include "wifi_monitor.h"
#include "gps_sync.h"  // <--- ADDED: GPS Support

static const char *TAG = "MAIN";

// --- Hardware Configuration ---
#if CONFIG_IDF_TARGET_ESP32C5
    #define RGB_LED_GPIO 27
#else
    // Fallback for other chips if you switch boards
    #define RGB_LED_GPIO 8
#endif

// --- LED State Machine ---
static led_strip_handle_t led_strip;
static bool wifi_connected = false;
static bool has_config = false;

typedef enum {
    LED_STATE_NO_CONFIG,   // Yellow Solid
    LED_STATE_WAITING,     // Blue Blink (Connecting)
    LED_STATE_CONNECTED,   // Green Solid (Connected to AP)
    LED_STATE_FAILED,      // Red Blink
    LED_STATE_MONITORING   // Blue Solid (Sniffing Air)
} led_state_t;

static led_state_t current_led_state = LED_STATE_NO_CONFIG;

static void rgb_led_init(void) {
    ESP_LOGI(TAG, "Initializing RGB LED on GPIO %d", RGB_LED_GPIO);
    led_strip_config_t strip_config = {
        .strip_gpio_num = RGB_LED_GPIO,
        .max_leds = 1,
    };
    led_strip_rmt_config_t rmt_config = {
        .resolution_hz = 10 * 1000 * 1000,
        .flags.with_dma = false,
    };
    ESP_ERROR_CHECK(led_strip_new_rmt_device(&strip_config, &rmt_config, &led_strip));
    led_strip_clear(led_strip);
}

static void set_led_color(uint8_t r, uint8_t g, uint8_t b) {
    led_strip_set_pixel(led_strip, 0, r, g, b);
    led_strip_refresh(led_strip);
}

static void led_task(void *arg) {
    int blink_state = 0;
    while(1) {
        switch(current_led_state) {
            case LED_STATE_NO_CONFIG:
                set_led_color(25, 25, 0); // Yellow (Dimmed)
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_WAITING:
                if (blink_state) set_led_color(0, 0, 50); // Blue
                else set_led_color(0, 0, 0);
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(500));
                break;

            case LED_STATE_CONNECTED:
                set_led_color(0, 25, 0); // Green
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_MONITORING:
                set_led_color(0, 0, 50); // Blue Solid
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_FAILED:
                if (blink_state) set_led_color(50, 0, 0); // Red
                else set_led_color(0, 0, 0);
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(200));
                break;
        }
    }
}

// --- GPS Logging Helper ---
// Replaces the old plain text log with your CSV + GPS Timestamp format
void log_collapse_event(float nav_duration_us, int rssi, int retry) {
    gps_timestamp_t ts = gps_get_timestamp();

    // Format: COLLAPSE,MonoMS,GpsMS,Synced,Duration,RSSI,Retry
    printf("COLLAPSE,%lld,%lld,%d,%.2f,%d,%d\n",
           ts.monotonic_ms,
           ts.gps_ms,
           ts.synced ? 1 : 0,
           nav_duration_us,
           rssi,
           retry);
}

// --- CSI Support ---------------------------------------------------
static bool s_csi_enabled = false;
static uint32_t s_csi_packet_count = 0;

static void csi_cb(void *ctx, wifi_csi_info_t *info) {
    csi_log_append_record(info);
    s_csi_packet_count++;
    if ((s_csi_packet_count % 100) == 0) {
        ESP_LOGI("CSI", "Captured %lu CSI packets", (unsigned long)s_csi_packet_count);
    }
}

static void wifi_enable_csi_once(void) {
    if (s_csi_enabled) return;

    vTaskDelay(pdMS_TO_TICKS(2000));
    wifi_csi_config_t csi_cfg;
    memset(&csi_cfg, 0, sizeof(csi_cfg));
    csi_cfg.enable = true; // C5 specific simple config

    ESP_LOGI("CSI", "Configuring CSI...");
    if (esp_wifi_set_csi_config(&csi_cfg) != ESP_OK) return;
    if (esp_wifi_set_csi_rx_cb(csi_cb, NULL) != ESP_OK) return;
    if (esp_wifi_set_csi(true) != ESP_OK) return;

    ESP_LOGI("CSI", "CSI enabled!");
    s_csi_enabled = true;
}

static void csi_dump_task(void *arg) {
    vTaskDelay(pdMS_TO_TICKS(20000)); // Dump after 20 seconds
    csi_log_dump_over_uart();
    vTaskDelete(NULL);
}

static void csi_init_task(void *arg) {
    wifi_enable_csi_once();
    vTaskDelete(NULL);
}

// --- WiFi Monitor Mode Support -------------------------------------
static bool s_monitor_enabled = false;
static uint32_t s_monitor_frame_count = 0;

// This is the core analysis function
static void monitor_frame_callback(const wifi_frame_info_t *frame,
                                   const uint8_t *payload,
                                   uint16_t len) {
    s_monitor_frame_count++;

    // 1. Check for Collapse (High NAV + Retry)
    if (frame->retry && frame->duration_id > 5000) {
        // USE GPS LOGGING HERE
        log_collapse_event((float)frame->duration_id, frame->rssi, frame->retry);
    }

    // 2. Also warn on extremely high NAV (blocking the channel)
    if (frame->duration_id > 30000) {
        ESP_LOGW("MONITOR", "⚠ VERY HIGH NAV: %u us", frame->duration_id);
    }
}

static void monitor_stats_task(void *arg) {
    while (1) {
        vTaskDelay(pdMS_TO_TICKS(10000)); // Every 10 seconds
        wifi_collapse_stats_t stats;
        if (wifi_monitor_get_stats(&stats) == ESP_OK) {
            ESP_LOGI("MONITOR", "--- Stats: %lu frames, Retry Rate: %.2f%%, Avg NAV: %u us ---",
                     stats.total_frames, stats.retry_rate, stats.avg_nav);

            if (wifi_monitor_is_collapsed()) {
                ESP_LOGW("MONITOR", "⚠⚠⚠ WiFi COLLAPSE DETECTED! ⚠⚠⚠");
            }
        }
    }
}

static void wifi_enable_monitor_mode(uint8_t channel) {
    if (s_monitor_enabled) return;

    ESP_LOGI("MONITOR", "Starting WiFi monitor mode on channel %d", channel);
    if (wifi_monitor_init(channel, monitor_frame_callback) != ESP_OK) return;
    if (wifi_monitor_start() != ESP_OK) return;

    s_monitor_enabled = true;
    current_led_state = LED_STATE_MONITORING;

    ESP_LOGI("MONITOR", "WiFi monitor started");
    xTaskCreate(monitor_stats_task, "monitor_stats", 4096, NULL, 5, NULL);
}

static void monitor_init_task(void *arg) {
    wifi_ap_record_t ap_info;
    // Try to sniff the same channel our AP is using
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        wifi_enable_monitor_mode(ap_info.primary);
    } else {
        wifi_enable_monitor_mode(6); // Default fallback
    }
    vTaskDelete(NULL);
}

// --- Event Handler (Connection Logic) ------------------------------
static void event_handler(void* arg, esp_event_base_t event_base,
                          int32_t event_id, void* event_data) {
    if (event_base == WIFI_EVENT) {
        if (event_id == WIFI_EVENT_STA_START) {
            if (has_config) current_led_state = LED_STATE_WAITING;
        }
        else if (event_id == WIFI_EVENT_STA_DISCONNECTED) {
            wifi_event_sta_disconnected_t* event = (wifi_event_sta_disconnected_t*) event_data;
            ESP_LOGW(TAG, "WiFi Disconnected (Reason: %d)", event->reason);
            if (!wifi_connected && has_config) current_led_state = LED_STATE_FAILED;
        }
    }
    else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
        ESP_LOGI(TAG, "Got IP: " IPSTR, IP2STR(&event->ip_info.ip));

        wifi_connected = true;
        current_led_state = LED_STATE_CONNECTED;

        // Sequence: 1. Start CSI, 2. Start Monitor, 3. Start Iperf
        xTaskCreate(csi_init_task, "csi_init", 4096, NULL, 5, NULL);

        vTaskDelay(pdMS_TO_TICKS(2000));
        xTaskCreate(monitor_init_task, "monitor_init", 4096, NULL, 5, NULL);

        vTaskDelay(pdMS_TO_TICKS(1000));
        iperf_cfg_t cfg;
        memset(&cfg, 0, sizeof(cfg));
        cfg.flag = IPERF_FLAG_SERVER | IPERF_FLAG_TCP;
        cfg.sport = 5001;
        iperf_start(&cfg);
        ESP_LOGI(TAG, "iperf TCP server started on port 5001");

        // Optional: Dump CSI data later
        xTaskCreate(csi_dump_task, "csi_dump_task", 4096, NULL, 5, NULL);
    }
}

// --- Main Application Entry ----------------------------------------
void app_main(void) {
    // 1. Initialize Non-Volatile Storage (needed for WiFi config)
    ESP_ERROR_CHECK(nvs_flash_init());

    // 2. Initialize Netif (TCP/IP stack)
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    // 3. Initialize Custom Logging & LED
    ESP_ERROR_CHECK(csi_log_init());
    rgb_led_init();
    xTaskCreate(led_task, "led_task", 4096, NULL, 5, NULL);

    // 4. Initialize GPS (The new addition!)
    // We do this EARLY so timestamps are ready when WiFi events happen
    ESP_LOGI(TAG, "Starting GPS Sync...");
    gps_sync_init(true); // true = Use GPS for system log timestamps

    // 5. Register WiFi Events
    ESP_ERROR_CHECK(esp_event_handler_instance_register(
                WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler, NULL, NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(
                IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler, NULL, NULL));

    // 6. Initialize WiFi Configuration
    wifi_cfg_init();

    if (wifi_cfg_apply_from_nvs()) {
        has_config = true;
        current_led_state = LED_STATE_WAITING;
        ESP_LOGI(TAG, "WiFi config loaded. Connecting...");
    } else {
        has_config = false;
        current_led_state = LED_STATE_NO_CONFIG;
        ESP_LOGI(TAG, "No WiFi config found. Yellow LED.");
        ESP_LOGI(TAG, "Use CLI 'wifi_config_set <ssid> <pass>' to configure.");
    }

    // 7. Loop forever (Logic is handled by tasks and events)
    while(1) {
        vTaskDelay(pdMS_TO_TICKS(1000));
        // Optional: Print GPS status occasionally
        if (!gps_is_synced()) {
             // ESP_LOGI(TAG, "Waiting for GPS lock...");
        }
    }
}