#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <inttypes.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <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"
#include "status_led.h"

static const char *TAG = "iperf";

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

#define RATE_CHECK_INTERVAL_US    500000
#define MIN_PACING_INTERVAL_US    100

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

static iperf_ctrl_t s_iperf_ctrl = {0};
static TaskHandle_t s_iperf_task_handle = NULL;
static iperf_cfg_t s_next_cfg;   // Holding area for the new config
static bool s_reload_req = false; // Flag to trigger internal restart

// Global Stats Tracker
static iperf_stats_t s_stats = {0};

// --- Session Persistence Variables ---
static int64_t s_session_start_time = 0;
static int64_t s_session_end_time = 0;
static uint64_t s_session_packets = 0;

// --- State Duration & Edge Counters ---
typedef enum {
    IPERF_STATE_IDLE = 0,
    IPERF_STATE_TX,
    IPERF_STATE_TX_SLOW,
    IPERF_STATE_TX_STALLED
} iperf_fsm_state_t;

static int64_t s_time_tx_us = 0;
static int64_t s_time_slow_us = 0;
static int64_t s_time_stalled_us = 0;

static uint32_t s_edge_tx = 0;
static uint32_t s_edge_slow = 0;
static uint32_t s_edge_stalled = 0;

static iperf_fsm_state_t s_current_fsm_state = IPERF_STATE_IDLE;

static esp_event_handler_instance_t instance_any_id;
static esp_event_handler_instance_t instance_got_ip;

// --- Helper: Pattern Initialization ---
// Fills buffer with 0-9 cyclic ASCII pattern (matches iperf2 "pattern" function)
static void iperf_pattern(uint8_t *buf, uint32_t len) {
    for (uint32_t i = 0; i < len; i++) {
        buf[i] = (i % 10) + '0';
    }
}

// --- Helper: Generate Client Header ---
// Modified to set all zeros except HEADER_SEQNO64B
static void iperf_generate_client_hdr(iperf_cfg_t *cfg, client_hdr_v1 *hdr) {
    // Zero out the entire structure
    memset(hdr, 0, sizeof(client_hdr_v1));

    // Set only the SEQNO64B flag (Server will detect 64-bit seqno in UDP header)
    hdr->flags = htonl(HEADER_SEQNO64B);
}

// ... [Existing Status Reporting & Event Handler Code] ...

void iperf_get_stats(iperf_stats_t *stats) {
    if (stats) {
        s_stats.config_pps = (s_iperf_ctrl.cfg.pacing_period_us > 0) ?
                             (1000000 / s_iperf_ctrl.cfg.pacing_period_us) : 0;
        *stats = s_stats;
    }
}

void iperf_print_status(void) {
    iperf_get_stats(&s_stats);

    // 1. Get Source IP
    char src_ip[32] = "0.0.0.0";
    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) {
             inet_ntop(AF_INET, &ip_info.ip, src_ip, sizeof(src_ip));
        }
    }

    // 2. Get Destination IP
    char dst_ip[32] = "0.0.0.0";
    struct in_addr daddr;
    daddr.s_addr = s_iperf_ctrl.cfg.dip;
    inet_ntop(AF_INET, &daddr, dst_ip, sizeof(dst_ip));

    float err = 0.0f;
    if (s_stats.running && s_stats.config_pps > 0) {
        int32_t diff = (int32_t)s_stats.config_pps - (int32_t)s_stats.actual_pps;
        err = (float)diff * 100.0f / (float)s_stats.config_pps;
    }

    // 3. Compute Session Bandwidth
    float avg_bw_mbps = 0.0f;
    if (s_session_start_time > 0) {
        int64_t end_t = (s_stats.running) ? esp_timer_get_time() : s_session_end_time;
        if (end_t > s_session_start_time) {
            double duration_sec = (double)(end_t - s_session_start_time) / 1000000.0;
            if (duration_sec > 0.001) {
                double total_bits = (double)s_session_packets * (double)s_iperf_ctrl.cfg.send_len * 8.0;
                avg_bw_mbps = (float)(total_bits / duration_sec / 1000000.0);
            }
        }
    }

    // 4. Calculate State Percentages
    double total_us = (double)(s_time_tx_us + s_time_slow_us + s_time_stalled_us);
    if (total_us < 1.0) total_us = 1.0;

    double pct_tx = ((double)s_time_tx_us / total_us) * 100.0;
    double pct_slow = ((double)s_time_slow_us / total_us) * 100.0;
    double pct_stalled = ((double)s_time_stalled_us / total_us) * 100.0;

    // Standard Stats
    printf("IPERF_STATUS: Src=%s, Dst=%s, Running=%d, Config=%" PRIu32 ", Actual=%" PRIu32 ", Err=%.1f%%, Pkts=%" PRIu64 ", AvgBW=%.2f Mbps\n",
           src_ip, dst_ip, s_stats.running, s_stats.config_pps, s_stats.actual_pps, err, s_session_packets, avg_bw_mbps);

    // New Format: Time + Percentage + Edges
    printf("IPERF_STATES: TX=%.2fs/%.1f%% (%lu), SLOW=%.2fs/%.1f%% (%lu), STALLED=%.2fs/%.1f%% (%lu)\n",
           (double)s_time_tx_us/1000000.0, pct_tx, (unsigned long)s_edge_tx,
           (double)s_time_slow_us/1000000.0, pct_slow, (unsigned long)s_edge_slow,
           (double)s_time_stalled_us/1000000.0, pct_stalled, (unsigned long)s_edge_stalled);
}

// --- Network Events ---
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 == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        xEventGroupSetBits(s_iperf_event_group, IPERF_IP_READY_BIT);
    }
    else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        xEventGroupClearBits(s_iperf_event_group, IPERF_IP_READY_BIT);
        status_led_set_state(LED_STATE_NO_CONFIG);
    }
}

static bool iperf_wait_for_ip(void) {
    if (!s_iperf_event_group) s_iperf_event_group = xEventGroupCreate();

    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 && ip_info.ip.addr != 0) {
            xEventGroupSetBits(s_iperf_event_group, IPERF_IP_READY_BIT);
        }
    }

    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_LOGI(TAG, "Waiting for IP...");
    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) {
        ESP_LOGW(TAG, "Stop requested while waiting for IP");
        return false;
    }
    ESP_LOGI(TAG, "IP Ready. Starting traffic.");
    return true;
}

static void trim_whitespace(char *str) {
    char *end = str + strlen(str) - 1;
    while(end > str && isspace((unsigned char)*end)) end--;
    *(end+1) = 0;
}

static void iperf_read_nvs_config(iperf_cfg_t *cfg) {
    nvs_handle_t my_handle;
    if (nvs_open("storage", NVS_READONLY, &my_handle) != ESP_OK) return;

    uint32_t val;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_PERIOD, &val) == ESP_OK) cfg->pacing_period_us = val;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_BURST, &val) == ESP_OK) cfg->burst_count = val;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_LEN, &val) == ESP_OK) cfg->send_len = val;
    if (nvs_get_u32(my_handle, NVS_KEY_IPERF_PORT, &val) == ESP_OK) cfg->dport = (uint16_t)val;

    size_t req;
    char buf[16];
    req = sizeof(buf);
    if (nvs_get_str(my_handle, NVS_KEY_IPERF_ROLE, buf, &req) == ESP_OK) {
        if (strcmp(buf, "SERVER") == 0) cfg->flag |= IPERF_FLAG_SERVER;
        else cfg->flag |= IPERF_FLAG_CLIENT;
    }

    req = sizeof(buf);
    if (nvs_get_str(my_handle, NVS_KEY_IPERF_PROTO, buf, &req) == ESP_OK) {
        if (strcmp(buf, "TCP") == 0) cfg->flag |= IPERF_FLAG_TCP;
        else cfg->flag |= IPERF_FLAG_UDP;
    }

    if (nvs_get_str(my_handle, NVS_KEY_IPERF_DST_IP, NULL, &req) == ESP_OK) {
        char *ip_str = malloc(req);
        if (ip_str) {
            nvs_get_str(my_handle, NVS_KEY_IPERF_DST_IP, ip_str, &req);
            trim_whitespace(ip_str);
            cfg->dip = inet_addr(ip_str);
            free(ip_str);
        }
    }
    nvs_close(my_handle);
}

void iperf_set_pps(uint32_t pps) {
    if (pps == 0) pps = 1;
    uint32_t period_us = 1000000 / pps;
    if (period_us < MIN_PACING_INTERVAL_US) period_us = MIN_PACING_INTERVAL_US;

    if (s_iperf_task_handle != NULL) {
        s_iperf_ctrl.cfg.pacing_period_us = period_us;
        printf("IPERF_PPS_UPDATED: %" PRIu32 "\n", pps);
    } else {
        s_iperf_ctrl.cfg.pacing_period_us = period_us;
    }
}

uint32_t iperf_get_pps(void) {
    if (s_iperf_ctrl.cfg.pacing_period_us == 0) return 0;
    return 1000000 / s_iperf_ctrl.cfg.pacing_period_us;
}


static esp_err_t iperf_start_udp_client(iperf_ctrl_t *ctrl) {
    if (!iperf_wait_for_ip()) {
        printf("IPERF_STOPPED\n");
        return ESP_OK;
    }

    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(ctrl->cfg.dport > 0 ? ctrl->cfg.dport : 5001);
    addr.sin_addr.s_addr = ctrl->cfg.dip;

    char ip_str[32];
    inet_ntop(AF_INET, &addr.sin_addr, ip_str, sizeof(ip_str));
    ESP_LOGI(TAG, "Client sending to %s:%d", ip_str, ntohs(addr.sin_port));

    int sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd < 0) {
        status_led_set_state(LED_STATE_FAILED);
        ESP_LOGE(TAG, "Socket creation failed: %d", errno);
        printf("IPERF_STOPPED\n");
        return ESP_FAIL;
    }

    status_led_set_state(LED_STATE_TRANSMITTING_SLOW);

    udp_datagram *udp_hdr = (udp_datagram *)ctrl->buffer;
    client_hdr_v1 *client_hdr = (client_hdr_v1 *)(ctrl->buffer + sizeof(udp_datagram));
    iperf_generate_client_hdr(&ctrl->cfg, client_hdr);

    s_stats.running = true;
    s_session_start_time = esp_timer_get_time();
    s_session_end_time = 0;
    s_session_packets = 0;

    // Reset FSM
    s_time_tx_us = 0; s_time_slow_us = 0; s_time_stalled_us = 0;
    s_edge_tx = 0; s_edge_slow = 0; s_edge_stalled = 0;
    s_current_fsm_state = IPERF_STATE_IDLE;

    printf("IPERF_STARTED\n");

    int64_t next_send_time = esp_timer_get_time();
    int64_t end_time = (ctrl->cfg.time == 0) ? INT64_MAX : esp_timer_get_time() + (int64_t)ctrl->cfg.time * 1000000LL;

    int64_t last_rate_check = esp_timer_get_time();
    uint32_t packets_since_check = 0;
    int64_t packet_id = 0;
    struct timespec ts;

    while (!ctrl->finish && esp_timer_get_time() < end_time) {
        int64_t now = esp_timer_get_time();
        int64_t wait = next_send_time - now;

        if (wait > 2000) vTaskDelay(pdMS_TO_TICKS(wait / 1000));
        else while (esp_timer_get_time() < next_send_time) taskYIELD();

        for (int k = 0; k < ctrl->cfg.burst_count; k++) {
            int64_t current_id = packet_id++;

            udp_hdr->id = htonl((uint32_t)(current_id & 0xFFFFFFFF));
            udp_hdr->id2 = htonl((uint32_t)((current_id >> 32) & 0xFFFFFFFF));

            clock_gettime(CLOCK_REALTIME, &ts);
            udp_hdr->tv_sec = htonl((uint32_t)ts.tv_sec);
            udp_hdr->tv_usec = htonl(ts.tv_nsec / 1000);

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

            if (sent > 0) {
                packets_since_check++;
                s_session_packets++;
            } else {
                // --- ROBUST FIX: Never Abort ---
                // If send fails (buffer full, routing issue, etc.), we just yield and retry next loop.
                // We do NOT goto exit.
                if (errno != 12) {
                    // Log rarely to avoid spamming serial
                    if ((packet_id % 100) == 0) {
                        ESP_LOGW(TAG, "Send error: %d (Ignored)", errno);
                    }
                }
                vTaskDelay(pdMS_TO_TICKS(10));
            }
        }

        now = esp_timer_get_time();
        if (now - last_rate_check > RATE_CHECK_INTERVAL_US) {
            uint32_t interval_us = (uint32_t)(now - last_rate_check);
            if (interval_us > 0) {
                s_stats.actual_pps = (uint32_t)((uint64_t)packets_since_check * 1000000 / interval_us);
                uint32_t config_pps = iperf_get_pps();
                uint32_t threshold = (config_pps * 3) / 4;
                iperf_fsm_state_t next_state;
                if (s_stats.actual_pps == 0) next_state = IPERF_STATE_TX_STALLED;
                else if (s_stats.actual_pps >= threshold) next_state = IPERF_STATE_TX;
                else next_state = IPERF_STATE_TX_SLOW;

                switch (next_state) {
                case IPERF_STATE_TX: s_time_tx_us += interval_us; break;
                case IPERF_STATE_TX_SLOW: s_time_slow_us += interval_us; break;
                case IPERF_STATE_TX_STALLED: s_time_stalled_us += interval_us; break;
                default: break;
                }
                if (next_state != s_current_fsm_state) {
                    switch (next_state) {
                    case IPERF_STATE_TX: s_edge_tx++; break;
                    case IPERF_STATE_TX_SLOW: s_edge_slow++; break;
                    case IPERF_STATE_TX_STALLED: s_edge_stalled++; break;
                    default: break;
                    }
                    s_current_fsm_state = next_state;
                }
                led_state_t led_target = (s_current_fsm_state == IPERF_STATE_TX) ? LED_STATE_TRANSMITTING : LED_STATE_TRANSMITTING_SLOW;
                if (status_led_get_state() != led_target) status_led_set_state(led_target);
            }
            last_rate_check = now;
            packets_since_check = 0;
        }
        next_send_time += ctrl->cfg.pacing_period_us;
    }

    udp_datagram *hdr = (udp_datagram *)ctrl->buffer;
    int64_t final_id = -packet_id;
    hdr->id = htonl((uint32_t)(final_id & 0xFFFFFFFF));
    hdr->id2 = htonl((uint32_t)((final_id >> 32) & 0xFFFFFFFF));

    clock_gettime(CLOCK_REALTIME, &ts);
    hdr->tv_sec = htonl((uint32_t)ts.tv_sec);
    hdr->tv_usec = htonl(ts.tv_nsec / 1000);
    for(int i=0; i<10; i++) {
        sendto(sockfd, ctrl->buffer, ctrl->cfg.send_len, 0, (struct sockaddr *)&addr, sizeof(addr));
        vTaskDelay(pdMS_TO_TICKS(2));
    }
    ESP_LOGI(TAG, "Sent termination packets (ID: %" PRId64 ")", final_id);

    close(sockfd);
    s_stats.running = false;
    s_session_end_time = esp_timer_get_time();
    s_stats.actual_pps = 0;
    status_led_set_state(LED_STATE_CONNECTED); // <--- This is your "Solid Green"
    printf("IPERF_STOPPED\n");
    return ESP_OK;
}

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

    do {
        s_reload_req = false;
        ctrl->finish = false;
        xEventGroupClearBits(s_iperf_event_group, IPERF_STOP_REQ_BIT);

        if (ctrl->cfg.flag & IPERF_FLAG_UDP && ctrl->cfg.flag & IPERF_FLAG_CLIENT) {
             iperf_start_udp_client(ctrl);
        }

        if (s_reload_req) {
            ESP_LOGI(TAG, "Hot reloading iperf task with new config...");
            ctrl->cfg = s_next_cfg;
            vTaskDelay(pdMS_TO_TICKS(100));
        }

    } while (s_reload_req);

    free(ctrl->buffer);
    s_iperf_task_handle = NULL;
    vTaskDelete(NULL);
}

void iperf_start(iperf_cfg_t *cfg) {
    iperf_cfg_t new_cfg = *cfg;
    iperf_read_nvs_config(&new_cfg);

    if (new_cfg.send_len == 0) new_cfg.send_len = 1470;
    if (new_cfg.pacing_period_us == 0) new_cfg.pacing_period_us = 10000;
    if (new_cfg.burst_count == 0) new_cfg.burst_count = 1;

    if (s_iperf_task_handle) {
        ESP_LOGI(TAG, "Task running. Staging hot reload.");
        s_next_cfg = new_cfg;
        s_reload_req = true;
        iperf_stop();
        printf("IPERF_RELOADING\n");
        return;
    }

    s_iperf_ctrl.cfg = new_cfg;
    s_iperf_ctrl.finish = false;

    if (s_iperf_ctrl.buffer == NULL) {
        s_iperf_ctrl.buffer_len = s_iperf_ctrl.cfg.send_len + 128;
        s_iperf_ctrl.buffer = calloc(1, s_iperf_ctrl.buffer_len);
    }

    // Initialize Buffer Pattern
    if (s_iperf_ctrl.buffer) {
        iperf_pattern(s_iperf_ctrl.buffer, s_iperf_ctrl.buffer_len);
    }

    if (s_iperf_event_group == NULL) {
        s_iperf_event_group = xEventGroupCreate();
    }

    xTaskCreate(iperf_task, "iperf", 4096, &s_iperf_ctrl, 5, &s_iperf_task_handle);
}

void iperf_stop(void) {
    if (s_iperf_task_handle) {
        s_iperf_ctrl.finish = true;
        if (s_iperf_event_group) xEventGroupSetBits(s_iperf_event_group, IPERF_STOP_REQ_BIT);
    } else {
        printf("IPERF_STOPPED\n");
    }
}