#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"

#include "iperf.h"
#include "wifi_cfg.h"
#include "csi_log.h"


static const char *TAG = "main";

#if CONFIG_IDF_TARGET_ESP32S3
    #define RGB_LED_GPIO 48
#elif CONFIG_IDF_TARGET_ESP32C5
    #define RGB_LED_GPIO 27
#elif CONFIG_IDF_TARGET_ESP32C6
    #define RGB_LED_GPIO 8
#elif CONFIG_IDF_TARGET_ESP32C3
    #define RGB_LED_GPIO 8
#elif CONFIG_IDF_TARGET_ESP32S2
    #define RGB_LED_GPIO 18
#elif CONFIG_IDF_TARGET_ESP32
    #define RGB_LED_GPIO 2
#else
    #error "Unsupported target - define RGB_LED_GPIO for your board"
#endif

static led_strip_handle_t led_strip;
static bool wifi_connected = false;
static bool has_config = false;

typedef enum {
    LED_STATE_NO_CONFIG,
    LED_STATE_WAITING,
    LED_STATE_CONNECTED,
    LED_STATE_FAILED
} led_state_t;

static led_state_t current_led_state = LED_STATE_NO_CONFIG;

static void rgb_led_init(void)
{
    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,
    };

    ESP_ERROR_CHECK(led_strip_new_rmt_device(&strip_config, &rmt_config, &led_strip));
    led_strip_clear(led_strip);

    ESP_LOGI(TAG, "WS2812 RGB LED initialized on GPIO %d", RGB_LED_GPIO);
}

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(255, 255, 0);
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_WAITING:
                if (blink_state) {
                    set_led_color(0, 0, 255);
                } else {
                    set_led_color(0, 0, 0);
                }
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_CONNECTED:
                set_led_color(0, 255, 0);
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_FAILED:
                if (blink_state) {
                    set_led_color(255, 0, 0);
                } else {
                    set_led_color(0, 0, 0);
                }
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(200));
                break;
        }
    }
}

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

static void csi_dump_task(void *arg) {
    vTaskDelay(pdMS_TO_TICKS(20000));
    csi_log_dump_over_uart();
    vTaskDelete(NULL);
}

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;
    }

    esp_err_t err;

    vTaskDelay(pdMS_TO_TICKS(2000));

    wifi_csi_config_t csi_cfg;

#if CONFIG_IDF_TARGET_ESP32C5
    memset(&csi_cfg, 0, sizeof(csi_cfg));
    csi_cfg.enable = true;

#elif CONFIG_IDF_TARGET_ESP32S3
    memset(&csi_cfg, 0, sizeof(csi_cfg));
    csi_cfg.lltf_en           = true;
    csi_cfg.htltf_en          = true;
    csi_cfg.stbc_htltf2_en    = false;
    csi_cfg.ltf_merge_en      = false;
    csi_cfg.channel_filter_en = false;
    csi_cfg.manu_scale        = false;
    csi_cfg.shift             = 0;

#else
#warning "CSI not supported for this target"
    return;
#endif

    ESP_LOGI("CSI", "Configuring CSI...");

    err = esp_wifi_set_csi_config(&csi_cfg);
    if (err != ESP_OK) {
        ESP_LOGE("CSI", "esp_wifi_set_csi_config failed: %s (0x%x)", esp_err_to_name(err), err);
        return;
    }
    ESP_LOGI("CSI", "CSI config OK");

    err = esp_wifi_set_csi_rx_cb(csi_cb, NULL);
    if (err != ESP_OK) {
        ESP_LOGE("CSI", "esp_wifi_set_csi_rx_cb failed: %s", esp_err_to_name(err));
        return;
    }
    ESP_LOGI("CSI", "CSI callback OK");

    err = esp_wifi_set_csi(true);
    if (err != ESP_OK) {
        ESP_LOGE("CSI", "esp_wifi_set_csi(true) failed: %s", esp_err_to_name(err));
        return;
    }
    ESP_LOGI("CSI", "CSI enabled!");

    s_csi_enabled = true;
}

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

static void event_handler(void* arg, esp_event_base_t event_base,
                          int32_t event_id, void* event_data)
{
    if (event_base == WIFI_EVENT) {
        switch (event_id) {
        case WIFI_EVENT_STA_START:
            ESP_LOGI(TAG, "WiFi started, attempting connection...");
            if (has_config) {
                current_led_state = LED_STATE_WAITING;
            }
            break;

        case WIFI_EVENT_STA_DISCONNECTED:
            wifi_event_sta_disconnected_t* event = (wifi_event_sta_disconnected_t*) event_data;

            // Get SSID for better error messages
            wifi_config_t wifi_cfg;
            const char *ssid = "unknown";
            if (esp_wifi_get_config(WIFI_IF_STA, &wifi_cfg) == ESP_OK) {
                ssid = (const char*)wifi_cfg.sta.ssid;
            }

            // Log disconnect with reason-specific messages
            switch (event->reason) {
                case 201: // WIFI_REASON_NO_AP_FOUND
                    ESP_LOGE(TAG, "WiFi disconnected (reason 201): NO AP FOUND");
                    ESP_LOGE(TAG, "  SSID attempted: '%s'", ssid);
                    ESP_LOGE(TAG, "  Verify AP is broadcasting and in range");
                    break;
                case 202: // WIFI_REASON_AUTH_FAIL
                    ESP_LOGE(TAG, "WiFi disconnected (reason 202): AUTH FAILED");
                    ESP_LOGE(TAG, "  SSID: '%s'", ssid);
                    ESP_LOGE(TAG, "  Check password is correct");
                    break;
                case 15: // WIFI_REASON_4WAY_HANDSHAKE_TIMEOUT
                    ESP_LOGE(TAG, "WiFi disconnected (reason 15): 4-WAY HANDSHAKE TIMEOUT");
                    ESP_LOGE(TAG, "  SSID: '%s'", ssid);
                    ESP_LOGE(TAG, "  Password may be incorrect");
                    break;
                case 2: // WIFI_REASON_AUTH_EXPIRE
                    ESP_LOGW(TAG, "WiFi disconnected (reason 2): AUTH EXPIRED");
                    ESP_LOGW(TAG, "  SSID: '%s' - will retry", ssid);
                    break;
                case 8: // WIFI_REASON_ASSOC_LEAVE
                    ESP_LOGW(TAG, "WiFi disconnected (reason 8): STATION LEFT");
                    ESP_LOGW(TAG, "  SSID: '%s' - normal disconnect", ssid);
                    break;
                default:
                    ESP_LOGW(TAG, "WiFi disconnected from '%s', reason: %d", ssid, event->reason);
                    break;
            }

            if (!wifi_connected && has_config) {
                current_led_state = LED_STATE_FAILED;
                ESP_LOGE(TAG, "WiFi connection FAILED - RED LED blinking");
            }
            break;
        }

    } 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 " gw:" IPSTR " netmask:" IPSTR,
                 IP2STR(&event->ip_info.ip),
                 IP2STR(&event->ip_info.gw),
                 IP2STR(&event->ip_info.netmask));

        wifi_connected = true;
        current_led_state = LED_STATE_CONNECTED;

        // Log connection details: SSID, band, channel, RSSI
        wifi_config_t wifi_cfg;
        wifi_ap_record_t ap_info;
        if (esp_wifi_get_config(WIFI_IF_STA, &wifi_cfg) == ESP_OK &&
            esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {

            const char *band_str = "Unknown";
            if (ap_info.primary >= 1 && ap_info.primary <= 14) {
                band_str = "2.4GHz";
            } else if (ap_info.primary >= 36) {
                band_str = "5GHz";
            }

            ESP_LOGI(TAG, "========================================");
            ESP_LOGI(TAG, "WiFi CONNECTED - BLUE LED solid");
            ESP_LOGI(TAG, "  SSID:    '%s'", wifi_cfg.sta.ssid);
            ESP_LOGI(TAG, "  Band:    %s", band_str);
            ESP_LOGI(TAG, "  Channel: %d", ap_info.primary);
            ESP_LOGI(TAG, "  RSSI:    %d dBm", ap_info.rssi);
            ESP_LOGI(TAG, "  BSSID:   %02x:%02x:%02x:%02x:%02x:%02x",
                     ap_info.bssid[0], ap_info.bssid[1], ap_info.bssid[2],
                     ap_info.bssid[3], ap_info.bssid[4], ap_info.bssid[5]);
            ESP_LOGI(TAG, "========================================");
        } else {
            ESP_LOGI(TAG, "WiFi CONNECTED - BLUE LED solid");
        }

        xTaskCreate(csi_init_task, "csi_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");

        xTaskCreate(csi_dump_task, "csi_dump_task", 4096, NULL, 5, NULL);
    }
}

void app_main(void) {
    ESP_ERROR_CHECK(nvs_flash_init());
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    ESP_ERROR_CHECK(csi_log_init());
    rgb_led_init();

    xTaskCreate(led_task, "led_task", 4096, NULL, 5, NULL);

    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));

    wifi_cfg_init();

    if (wifi_cfg_apply_from_nvs()) {
        has_config = true;
        current_led_state = LED_STATE_WAITING;
        ESP_LOGI(TAG, "WiFi config loaded from NVS");
    } else {
        has_config = false;
        current_led_state = LED_STATE_NO_CONFIG;
        ESP_LOGI(TAG, "No WiFi config - YELLOW LED");
    }

    ESP_LOGI(TAG, "LED Status:");
    ESP_LOGI(TAG, "  YELLOW solid    = NO CONFIG (send CFG/END)");
    ESP_LOGI(TAG, "  BLUE slow blink = Connecting");
    ESP_LOGI(TAG, "  BLUE solid      = Connected ✓");
    ESP_LOGI(TAG, "  RED fast blink  = Failed ✗");

    while(1) {
        vTaskDelay(pdMS_TO_TICKS(1000));
    }
}