#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_log.h"
#include "esp_err.h"
#include "esp_timer.h"
#include "nvs_flash.h"
#include "nvs.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "esp_wifi.h"
#include "iperf.h"

static const char *TAG = "iperf";

// --- LED STATE MANAGEMENT ---
typedef enum {
    LED_OFF,
    LED_BLUE_SOLID,     // Monitor Mode
    LED_RED_FLASH,      // No WiFi
    LED_AMBER_SOLID,    // Connected, No IP
    LED_GREEN_SOLID,    // Got IP / Ready
    LED_PURPLE_SOLID,   // Transmitting
    LED_PURPLE_FLASH    // Socket Error
} led_state_t;

static led_state_t s_led_state = LED_RED_FLASH;

// --- Helper: Set Physical LED ---
static void iperf_set_physical_led(uint8_t r, uint8_t g, uint8_t b) {
    // Implement hardware specific LED driver here
}

// --- LED Task ---
static void status_led_task(void *arg) {
    bool toggle = false;
    while (1) {
        switch (s_led_state) {
            case LED_BLUE_SOLID:   iperf_set_physical_led(0, 0, 64); vTaskDelay(pdMS_TO_TICKS(500)); break;
            case LED_RED_FLASH:    iperf_set_physical_led(toggle ? 64 : 0, 0, 0); vTaskDelay(pdMS_TO_TICKS(250)); toggle = !toggle; break;
            case LED_AMBER_SOLID:  iperf_set_physical_led(32, 16, 0); vTaskDelay(pdMS_TO_TICKS(500)); break;
            case LED_GREEN_SOLID:  iperf_set_physical_led(0, 64, 0); vTaskDelay(pdMS_TO_TICKS(500)); break;
            case LED_PURPLE_SOLID: iperf_set_physical_led(64, 0, 64); vTaskDelay(pdMS_TO_TICKS(200)); break;
            case LED_PURPLE_FLASH: iperf_set_physical_led(toggle ? 64 : 0, 0, 64); vTaskDelay(pdMS_TO_TICKS(250)); toggle = !toggle; break;
            default:               iperf_set_physical_led(0, 0, 0); vTaskDelay(pdMS_TO_TICKS(500)); break;
        }
    }
}

// --- Synchronization ---
static EventGroupHandle_t s_iperf_event_group = NULL;
#define IPERF_IP_READY_BIT   (1 << 0)
#define IPERF_STOP_REQ_BIT   (1 << 1)

typedef struct {
    iperf_cfg_t cfg;
    bool finish;
    uint32_t total_len;
    uint32_t buffer_len;
    uint8_t *buffer;
    uint32_t sockfd;
} iperf_ctrl_t;

static iperf_ctrl_t s_iperf_ctrl = {0};
static TaskHandle_t s_iperf_task_handle = NULL;
static esp_event_handler_instance_t instance_any_id;
static esp_event_handler_instance_t instance_got_ip;

// --- Network Event Handler ---
static void iperf_network_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data)
{
    if (s_iperf_event_group == NULL) return;

    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_CONNECTED) {
        s_led_state = LED_AMBER_SOLID;
    }
    else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        xEventGroupSetBits(s_iperf_event_group, IPERF_IP_READY_BIT);
        if (s_led_state != LED_PURPLE_SOLID && s_led_state != LED_PURPLE_FLASH) {
            s_led_state = LED_GREEN_SOLID;
        }
    }
    else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        xEventGroupClearBits(s_iperf_event_group, IPERF_IP_READY_BIT);
        s_led_state = LED_RED_FLASH;
    }
}

// --- Wait for IP ---
static bool iperf_wait_for_ip(void) {
    if (!s_iperf_event_group) s_iperf_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &iperf_network_event_handler, NULL, &instance_any_id));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &iperf_network_event_handler, NULL, &instance_got_ip));

    esp_netif_t *netif = esp_netif_get_handle_from_ifkey("WIFI_STA_DEF");
    if (netif) {
        esp_netif_ip_info_t ip_info;
        if (esp_netif_get_ip_info(netif, &ip_info) == ESP_OK) {
            if (ip_info.ip.addr != 0) {
                xEventGroupSetBits(s_iperf_event_group, IPERF_IP_READY_BIT);
                s_led_state = LED_GREEN_SOLID;
            } else {
                wifi_ap_record_t ap_info;
                if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) s_led_state = LED_AMBER_SOLID;
            }
        }
    }

    ESP_LOGI(TAG, "Waiting for IP address...");
    EventBits_t bits = xEventGroupWaitBits(s_iperf_event_group, IPERF_IP_READY_BIT | IPERF_STOP_REQ_BIT, pdFALSE, pdFALSE, portMAX_DELAY);

    esp_event_handler_instance_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, instance_any_id);
    esp_event_handler_instance_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, instance_got_ip);

    if (bits & IPERF_STOP_REQ_BIT) return false;
    return true;
}

// --- Read NVS ---
static void iperf_read_nvs_config(iperf_cfg_t *cfg) {
    nvs_handle_t my_handle;
    esp_err_t err = nvs_open("storage", NVS_READONLY, &my_handle);
    if (err != ESP_OK) return;
    size_t required_size;
    uint32_t val = 0;

    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_RATE, &val) == ESP_OK && val > 0) cfg->bw_lim = val;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_BURST, &val) == ESP_OK && val > 0) cfg->burst_count = val; else cfg->burst_count = 1;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_LEN, &val) == ESP_OK && val > 0) cfg->send_len = val; else cfg->send_len = IPERF_UDP_TX_LEN;

    if (nvs_get_str(my_handle, NVS_KEY_IPERF_DST_IP, NULL, &required_size) == ESP_OK) {
        char *ip_str = malloc(required_size);
        if (ip_str) { nvs_get_str(my_handle, NVS_KEY_IPERF_DST_IP, ip_str, &required_size); cfg->dip = inet_addr(ip_str); free(ip_str); }
    }
    if (nvs_get_str(my_handle, NVS_KEY_IPERF_ROLE, NULL, &required_size) == ESP_OK) {
        char *role = malloc(required_size);
        if (role) {
            nvs_get_str(my_handle, NVS_KEY_IPERF_ROLE, role, &required_size);
            if (strcmp(role, "SERVER") == 0) { cfg->flag &= ~IPERF_FLAG_CLIENT; cfg->flag |= IPERF_FLAG_SERVER; }
            else { cfg->flag &= ~IPERF_FLAG_SERVER; cfg->flag |= IPERF_FLAG_CLIENT; }
            free(role);
        }
    }
    if (nvs_get_str(my_handle, NVS_KEY_IPERF_PROTO, NULL, &required_size) == ESP_OK) {
        char *proto = malloc(required_size);
        if (proto) {
            nvs_get_str(my_handle, NVS_KEY_IPERF_PROTO, proto, &required_size);
            if (strcmp(proto, "TCP") == 0) { cfg->flag &= ~IPERF_FLAG_UDP; cfg->flag |= IPERF_FLAG_TCP; }
            else { cfg->flag &= ~IPERF_FLAG_TCP; cfg->flag |= IPERF_FLAG_UDP; }
            free(proto);
        }
    }
    nvs_close(my_handle);
}

// --- Stubbed / Unused Functions (Marked to silence warnings) ---
static void __attribute__((unused)) socket_send(int sockfd, const uint8_t *buffer, int len) {
    // Stub
}
static int __attribute__((unused)) socket_recv(int sockfd, uint8_t *buffer, int len, TickType_t timeout_ticks) {
    return 0; // Stub
}
static esp_err_t iperf_start_tcp_server(iperf_ctrl_t *ctrl) {
    ESP_LOGW(TAG, "TCP Server not implemented");
    return ESP_FAIL;
}
static esp_err_t iperf_start_tcp_client(iperf_ctrl_t *ctrl) {
    ESP_LOGW(TAG, "TCP Client not implemented");
    return ESP_FAIL;
}
static esp_err_t iperf_start_udp_server(iperf_ctrl_t *ctrl) {
    ESP_LOGW(TAG, "UDP Server not implemented");
    return ESP_FAIL;
}

// -----------------------------------------------------------------------------
// MAIN UDP CLIENT
// -----------------------------------------------------------------------------
static esp_err_t iperf_start_udp_client(iperf_ctrl_t *ctrl)
{
    if (!iperf_wait_for_ip()) return ESP_FAIL;

    struct sockaddr_in addr;
    int sockfd;
    struct timeval tv;

    sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        s_led_state = LED_RED_FLASH;
        return ESP_FAIL;
    }

    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.dport);
    addr.sin_addr.s_addr = ctrl->cfg.dip;

    uint32_t burst_count = ctrl->cfg.burst_count ? ctrl->cfg.burst_count : 1;
    uint32_t payload_len = ctrl->cfg.send_len ? ctrl->cfg.send_len : IPERF_UDP_TX_LEN;
    double target_bandwidth_mbps = (double)ctrl->cfg.bw_lim;
    if (target_bandwidth_mbps <= 0) target_bandwidth_mbps = 1.0;

    double target_bps = target_bandwidth_mbps * 1000000.0;
    double total_pps = target_bps / (payload_len * 8.0);
    double bursts_per_sec = total_pps / (double)burst_count;
    double pacing_interval_us = 1000000.0 / bursts_per_sec;

    s_led_state = LED_PURPLE_SOLID; // Transmitting

    uint64_t total_len = 0;
    uint32_t packet_count = 0;
    int64_t start_time_us = esp_timer_get_time();
    int64_t next_send_time = start_time_us;
    int64_t end_time_us = start_time_us + ((int64_t)ctrl->cfg.time * 1000000LL);
    double interval_accum = 0.0;

    while (!ctrl->finish && esp_timer_get_time() < end_time_us) {
        int64_t current_time = esp_timer_get_time();

        if (current_time >= next_send_time) {
            for (int k = 0; k < burst_count; k++) {
                udp_datagram *header = (udp_datagram *)ctrl->buffer;
                gettimeofday(&tv, NULL);
                header->id = htonl(packet_count);
                header->tv_sec = htonl(tv.tv_sec);
                header->tv_usec = htonl(tv.tv_usec);
                header->id2 = 0;

                if (packet_count == 0) {
                    client_hdr_v1 *client_hdr = (client_hdr_v1 *)(ctrl->buffer + sizeof(udp_datagram));
                    client_hdr->flags = htonl(HEADER_VERSION1);
                    client_hdr->numThreads = htonl(1);
                    client_hdr->mPort = htonl(ctrl->cfg.dport);
                    client_hdr->mBufLen = htonl(payload_len);
                    client_hdr->mWinBand = htonl((int)target_bps);
                    client_hdr->mAmount = htonl(-(int)(ctrl->cfg.time * 100));
                }

                int send_len = sendto(sockfd, ctrl->buffer, payload_len, 0, (struct sockaddr *)&addr, sizeof(addr));

                if (send_len > 0) {
                    total_len += send_len;
                    packet_count++;
                } else {
                    ESP_LOGE(TAG, "UDP send failed: %d", errno);
                    s_led_state = LED_PURPLE_FLASH;
                    goto exit_client;
                }
            }
            interval_accum += pacing_interval_us;
            int64_t steps = (int64_t)interval_accum;
            if (steps > 0) { next_send_time += steps; interval_accum -= steps; }
            if (esp_timer_get_time() > next_send_time + 4000) next_send_time = esp_timer_get_time() + (int64_t)pacing_interval_us;
        } else {
            int64_t wait = next_send_time - current_time;
            if (wait > 2000) vTaskDelay(pdMS_TO_TICKS(wait/1000));
        }
    }

exit_client:
    if (s_led_state != LED_PURPLE_FLASH) s_led_state = LED_GREEN_SOLID;
    close(sockfd);
    return ESP_OK;
}

static void iperf_task(void *arg)
{
    iperf_ctrl_t *ctrl = (iperf_ctrl_t *)arg;

    if (ctrl->cfg.flag & IPERF_FLAG_TCP) {
        if (ctrl->cfg.flag & IPERF_FLAG_SERVER) iperf_start_tcp_server(ctrl);
        else iperf_start_tcp_client(ctrl);
    } else {
        if (ctrl->cfg.flag & IPERF_FLAG_SERVER) iperf_start_udp_server(ctrl);
        else iperf_start_udp_client(ctrl);
    }

    if (ctrl->buffer) { free(ctrl->buffer); ctrl->buffer = NULL; }
    if (s_iperf_event_group) { vEventGroupDelete(s_iperf_event_group); s_iperf_event_group = NULL; }

    s_iperf_task_handle = NULL;
    vTaskDelete(NULL);
}

void iperf_start(iperf_cfg_t *cfg)
{
    static bool led_task_started = false;
    if (!led_task_started) {
        xTaskCreate(status_led_task, "status_led", 2048, NULL, 1, NULL);
        led_task_started = true;
    }

    nvs_handle_t my_handle;
    uint8_t enabled = 1;
    if (nvs_open("storage", NVS_READONLY, &my_handle) == ESP_OK) {
        nvs_get_u8(my_handle, NVS_KEY_IPERF_ENABLE, &enabled);
        size_t req;
        if (nvs_get_str(my_handle, "mode", NULL, &req) == ESP_OK) {
             char m[10]; nvs_get_str(my_handle, "mode", m, &req);
             if (strcmp(m, "MONITOR") == 0) s_led_state = LED_BLUE_SOLID;
        }
        nvs_close(my_handle);
    }

    if (enabled == 0) return;
    if (s_iperf_task_handle != NULL) return;

    memcpy(&s_iperf_ctrl.cfg, cfg, sizeof(iperf_cfg_t));
    iperf_read_nvs_config(&s_iperf_ctrl.cfg);
    s_iperf_ctrl.finish = false;

    s_iperf_ctrl.buffer_len = 2048;
    s_iperf_ctrl.buffer = malloc(s_iperf_ctrl.buffer_len);
    memset(s_iperf_ctrl.buffer, 0, s_iperf_ctrl.buffer_len);

    s_iperf_event_group = xEventGroupCreate();
    xTaskCreate(iperf_task, "iperf", 4096, &s_iperf_ctrl, IPERF_TRAFFIC_TASK_PRIORITY, &s_iperf_task_handle);
}

void iperf_stop(void) {
    if (s_iperf_task_handle != NULL) {
        s_iperf_ctrl.finish = true;
        if (s_iperf_event_group) xEventGroupSetBits(s_iperf_event_group, IPERF_STOP_REQ_BIT);
    }
}