#include "wifi_controller.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "inttypes.h"

// Dependencies
#include "iperf.h"
#include "status_led.h"
#include "wifi_monitor.h"
#include "gps_sync.h"
#include "wifi_cfg.h"

// 1. GUARDED INCLUDE
#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
#include "csi_manager.h"
#endif

static const char *TAG = "WIFI_CTL";

static wifi_ctl_mode_t s_current_mode = WIFI_CTL_MODE_STA;
static uint8_t s_monitor_channel = 6;
static bool s_monitor_enabled = false;
static uint32_t s_monitor_frame_count = 0;
static TaskHandle_t s_monitor_stats_task_handle = NULL;
static uint8_t s_monitor_channel_staging = 6; //RAM Staging Variable

// --- Helper: Log Collapse Events ---
static void log_collapse_event(float nav_duration_us, int rssi, int retry) {
    gps_timestamp_t ts = gps_get_timestamp();
    // CSV Format: COLLAPSE,MonoMS,GpsMS,Synced,Duration,RSSI,Retry
    printf("COLLAPSE,%" PRIi64 ",%" PRIi64 ",%d,%.2f,%d,%d\n",
           ts.monotonic_ms,
           ts.gps_ms,
           ts.synced ? 1 : 0,
           nav_duration_us,
           rssi,
           retry);
}

// --- Monitor Callbacks & Tasks ---

static void monitor_frame_callback(const wifi_frame_info_t *frame, const uint8_t *payload, uint16_t len) {
    s_monitor_frame_count++;

    // Check for Collapse conditions (High NAV + Retry)
    if (frame->retry && frame->duration_id > 5000) {
        log_collapse_event((float)frame->duration_id, frame->rssi, frame->retry);
    }

    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));
        wifi_collapse_stats_t stats;
        if (wifi_monitor_get_stats(&stats) == ESP_OK) {
            ESP_LOGI("MONITOR", "--- Stats: %lu frames, Retry: %.2f%%, Avg NAV: %u us ---",
                     (unsigned long)stats.total_frames, stats.retry_rate, stats.avg_nav);

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

static void auto_monitor_task_func(void *arg) {
    uint8_t channel = (uint8_t)(uintptr_t)arg;

    ESP_LOGI(TAG, "Waiting for WiFi connection before switching to monitor mode...");
    // Wait until LED indicates connected
    while (status_led_get_state() != LED_STATE_CONNECTED) {
        vTaskDelay(pdMS_TO_TICKS(500));
    }

    ESP_LOGI(TAG, "WiFi connected, waiting for GPS sync (2s)...");
    vTaskDelay(pdMS_TO_TICKS(2000));

    ESP_LOGI(TAG, "Auto-switching to MONITOR mode on channel %d...", channel);
    wifi_ctl_switch_to_monitor(channel, WIFI_BW_HT20);
    vTaskDelete(NULL);
}

// --- API Implementation ---
void wifi_ctl_init(void) {
    s_current_mode = WIFI_CTL_MODE_STA;
    s_monitor_enabled = false;
    s_monitor_frame_count = 0;

    // Load Initial Staging from NVS
    char mode_ignored[16];
    wifi_cfg_get_mode(mode_ignored, &s_monitor_channel_staging);
    if (s_monitor_channel_staging == 0) s_monitor_channel_staging = 6;
}

void wifi_ctl_param_set_monitor_channel(uint8_t channel) {
    if (channel >= 1 && channel <= 14) {
        s_monitor_channel_staging = channel;
    }
}

uint8_t wifi_ctl_param_get_monitor_channel(void) {
    return s_monitor_channel_staging;
}

bool wifi_ctl_param_save(void) {
    bool changed = wifi_cfg_set_monitor_channel(s_monitor_channel_staging);
    if (changed) {
        ESP_LOGI(TAG, "Monitor channel (%d) saved to NVS", s_monitor_channel_staging);
    }
    return changed;
}

void wifi_ctl_param_reload(void) {
    char mode_ignored[16];
    uint8_t ch = 0;
    wifi_cfg_get_mode(mode_ignored, &ch);
    if (ch > 0) s_monitor_channel_staging = ch;
    ESP_LOGI(TAG, "Reloaded monitor channel: %d", s_monitor_channel_staging);
}

bool wifi_ctl_param_is_unsaved(void) {
    return wifi_cfg_monitor_channel_is_unsaved(s_monitor_channel_staging);
}
esp_err_t wifi_ctl_switch_to_monitor(uint8_t channel_override, wifi_bandwidth_t bandwidth) {
    // If override is 0, use Staging
    uint8_t channel = (channel_override > 0) ? channel_override : s_monitor_channel_staging;

    if (s_current_mode == WIFI_CTL_MODE_MONITOR && s_monitor_channel == channel) {
        ESP_LOGW(TAG, "Already in monitor mode");
        return ESP_OK;
    }

    // Monitor mode typically requires 20MHz
    if (bandwidth != WIFI_BW_HT20) {
        ESP_LOGW(TAG, "Forcing bandwidth to 20MHz for monitor mode");
        bandwidth = WIFI_BW_HT20;
    }

    ESP_LOGI(TAG, "Switching to MONITOR MODE (Ch %d)", channel);

    // 1. Stop high-level apps
    iperf_stop();
    vTaskDelay(pdMS_TO_TICKS(500));

    // 2. Disable CSI (hardware conflict)
    // 2. GUARDED CALL
#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
    csi_mgr_disable();
#endif

    // 3. Teardown Station
    esp_wifi_disconnect();
    esp_wifi_stop();
    vTaskDelay(pdMS_TO_TICKS(500));

    // 4. Re-init in NULL/Promiscuous Mode
    esp_wifi_set_mode(WIFI_MODE_NULL);

    if (wifi_monitor_init(channel, monitor_frame_callback) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to init monitor mode");
        return ESP_FAIL;
    }

    esp_wifi_set_bandwidth(WIFI_IF_STA, bandwidth);

    if (wifi_monitor_start() != ESP_OK) {
        ESP_LOGE(TAG, "Failed to start monitor mode");
        return ESP_FAIL;
    }

    // 5. Update State
    s_monitor_enabled = true;
    s_current_mode = WIFI_CTL_MODE_MONITOR;
    s_monitor_channel = channel;
    status_led_set_state(LED_STATE_MONITORING);

    if (s_monitor_stats_task_handle == NULL) {
        xTaskCreate(monitor_stats_task, "monitor_stats", 4096, NULL, 5, &s_monitor_stats_task_handle);
    }

    return ESP_OK;
}

esp_err_t wifi_ctl_switch_to_sta(wifi_band_mode_t band_mode) {
    if (s_current_mode == WIFI_CTL_MODE_STA) {
        ESP_LOGW(TAG, "Already in STA mode");
        return ESP_OK;
    }

    ESP_LOGI(TAG, "Switching to STA MODE");

    // 1. Stop Monitor Tasks
    if (s_monitor_stats_task_handle != NULL) {
        vTaskDelete(s_monitor_stats_task_handle);
        s_monitor_stats_task_handle = NULL;
    }

    // 2. Stop Monitor Driver
    if (s_monitor_enabled) {
        wifi_monitor_stop();
        s_monitor_enabled = false;
        vTaskDelay(pdMS_TO_TICKS(500));
    }

    // 3. Re-enable Station Mode
    esp_wifi_set_mode(WIFI_MODE_STA);
    vTaskDelay(pdMS_TO_TICKS(500));

    // 4. Configure & Connect
    wifi_config_t wifi_config;
    esp_wifi_get_config(WIFI_IF_STA, &wifi_config);
    wifi_config.sta.channel = 0; // Auto channel scan

    esp_wifi_set_config(WIFI_IF_STA, &wifi_config);
    esp_wifi_start();
    vTaskDelay(pdMS_TO_TICKS(500));
    esp_wifi_connect();

    // 5. Update State
    s_current_mode = WIFI_CTL_MODE_STA;
    status_led_set_state(LED_STATE_WAITING);

    return ESP_OK;
}

void wifi_ctl_auto_monitor_start(uint8_t channel) {
    xTaskCreate(auto_monitor_task_func, "auto_monitor", 4096, (void*)(uintptr_t)channel, 5, NULL);
}

wifi_ctl_mode_t wifi_ctl_get_mode(void) {
    return s_current_mode;
}

uint8_t wifi_ctl_get_monitor_channel(void) {
    return s_monitor_channel;
}

uint32_t wifi_ctl_get_monitor_frame_count(void) {
    return s_monitor_frame_count;
}