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

static const char *TAG = "main";

// RGB LED GPIO pin configuration
#if CONFIG_IDF_TARGET_ESP32S3
    // ESP32-S3-DevKitC-1 v1.0 uses GPIO48, v1.1 uses GPIO38
    // Default to GPIO48 for backward compatibility, but you may need GPIO38
    #define RGB_LED_GPIO 48  // Change to 38 if you have v1.1 board
#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;

// WiFi connection status
static bool wifi_connected = false;
static bool has_config = false;

// LED states
typedef enum {
    LED_STATE_NO_CONFIG,    // Yellow - no WiFi config in NVS
    LED_STATE_WAITING,      // Blue slow blink - trying to connect
    LED_STATE_CONNECTED,    // Blue solid - connected successfully
    LED_STATE_FAILED        // Red fast blink - connection 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, // 10MHz
    };

    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:
                // Yellow solid - no WiFi config
                set_led_color(255, 255, 0);  // Red + Green = Yellow
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_WAITING:
                // Blue slow blink - waiting for WiFi
                if (blink_state) {
                    set_led_color(0, 0, 255);  // Blue
                } else {
                    set_led_color(0, 0, 0);    // Off
                }
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_CONNECTED:
                // Blue solid - connected successfully
                set_led_color(0, 255, 0);  // Green
                vTaskDelay(pdMS_TO_TICKS(1000));
                break;

            case LED_STATE_FAILED:
                // Red fast blink - connection failed
                if (blink_state) {
                    set_led_color(255, 0, 0);  // Red
                } else {
                    set_led_color(0, 0, 0);    // Off
                }
                blink_state = !blink_state;
                vTaskDelay(pdMS_TO_TICKS(200));
                break;
        }
    }
}

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;
                ESP_LOGW(TAG, "WiFi disconnected, reason: %d", event->reason);

                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;
        ESP_LOGI(TAG, "WiFi CONNECTED - BLUE LED solid");

        // Auto-start iperf server
        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");
    }
}

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

    // Initialize RGB LED
    rgb_led_init();

    // Start LED task
    xTaskCreate(led_task, "led_task", 4096, NULL, 5, NULL);

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

    // Initialize WiFi config
    wifi_cfg_init();

    // Try to load config from NVS
    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));
    }
}