#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdatomic.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "nvs.h"
#include "esp_netif.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_check.h"

#include "wifi_cfg.h"
#include "cmd_transport.h"

#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
#include "csi_manager.h"
#endif

static const char *TAG = "wifi_cfg";
static esp_netif_t *sta_netif = NULL;
static bool cfg_dhcp = true;

// --- NVS Helper ---
static esp_err_t nvs_set_str2(nvs_handle_t h, const char *key, const char *val){
    return val ? nvs_set_str(h, key, val) : nvs_erase_key(h, key);
}

// --- NVS Save Functions ---

// 1. Save Network Settings (Namespace: "netcfg")
static void save_net_cfg(const char* ssid, const char* pass, const char* ip, const char* mask, const char* gw, bool dhcp, const char* band, const char* bw, const char* powersave, const char* mode, uint8_t mon_ch){
    nvs_handle_t h;
    if (nvs_open("netcfg", NVS_READWRITE, &h) != ESP_OK) return;

    if (ssid) nvs_set_str2(h, "ssid", ssid);
    if (pass) nvs_set_str2(h, "pass", pass);
    if (ip)   nvs_set_str2(h, "ip",   ip);
    if (mask) nvs_set_str2(h, "mask", mask);
    if (gw)   nvs_set_str2(h, "gw",   gw);
    if (band) nvs_set_str2(h, "band", band);
    if (bw)   nvs_set_str2(h, "bw",   bw);
    if (powersave) nvs_set_str2(h, "powersave", powersave);
    if (mode) nvs_set_str2(h, "mode", mode);
    nvs_set_u8(h, "mon_ch", mon_ch);
    nvs_set_u8(h, "dhcp", dhcp ? 1 : 0);

    nvs_commit(h);
    nvs_close(h);
    cfg_dhcp = dhcp;
    ESP_LOGI(TAG, "Net Config Saved: SSID=%s IP=%s DHCP=%d", ssid?ssid:"", ip?ip:"", dhcp);
}

// 2. Save Iperf Settings (Namespace: "storage") -> Matches iperf.c keys
static void save_iperf_cfg(const char* dst_ip, const char* role, const char* proto, uint32_t period, uint32_t burst, uint32_t len, uint32_t port, bool enable){
    nvs_handle_t h;
    if (nvs_open("storage", NVS_READWRITE, &h) != ESP_OK) return;

    // Note: Keys must match what iperf.c reads (NVS_KEY_IPERF_DST_IP etc)
    if (dst_ip && dst_ip[0]) nvs_set_str(h, "iperf_dst_ip", dst_ip);
    if (role && role[0])     nvs_set_str(h, "iperf_role",   role);
    if (proto && proto[0])   nvs_set_str(h, "iperf_proto",  proto);

    nvs_set_u32(h, "iperf_period", period);
    nvs_set_u32(h, "iperf_burst",  burst);
    nvs_set_u32(h, "iperf_len",    len);
    nvs_set_u32(h, "iperf_port",   port);
    nvs_set_u8(h,  "iperf_enabled", enable ? 1 : 0);

    nvs_commit(h);
    nvs_close(h);
    ESP_LOGI(TAG, "Iperf Config Saved: Target=%s Role=%s Period=%lu", dst_ip?dst_ip:"", role?role:"", (unsigned long)period);
}

// --- Load Logic (Network Only - Iperf loads itself) ---
static bool load_cfg(char* ssid, size_t ssz, char* pass, size_t psz,
                     char* ip, size_t isz, char* mask, size_t msz, char* gw, size_t gsz,
                     char* band, size_t bsz, char* bw, size_t bwsz, char* powersave, size_t pssz,
                     char* mode, size_t modesz, uint8_t* mon_ch, bool* dhcp){
    nvs_handle_t h;
    if (nvs_open("netcfg", NVS_READONLY, &h) != ESP_OK) return false;
    size_t len;
    esp_err_t e;
    if ((e = nvs_get_str(h, "ssid", NULL, &len)) != ESP_OK){ nvs_close(h); return false; }
    if (len >= ssz){ nvs_close(h); return false; }
    nvs_get_str(h, "ssid", ssid, &len);
    len = psz; e = nvs_get_str(h, "pass", pass, &len); if (e!=ESP_OK) pass[0]=0;
    len = isz; e = nvs_get_str(h, "ip", ip, &len);     if (e!=ESP_OK) ip[0]=0;
    len = msz; e = nvs_get_str(h, "mask", mask, &len); if (e!=ESP_OK) mask[0]=0;
    len = gsz; e = nvs_get_str(h, "gw", gw, &len);     if (e!=ESP_OK) gw[0]=0;
    len = bsz; e = nvs_get_str(h, "band", band, &len); if (e!=ESP_OK) strcpy(band, "2.4G");
    len = bwsz; e = nvs_get_str(h, "bw", bw, &len);    if (e!=ESP_OK) strcpy(bw, "HT20");
    len = pssz; e = nvs_get_str(h, "powersave", powersave, &len); if (e!=ESP_OK) strcpy(powersave, "NONE");
    len = modesz; e = nvs_get_str(h, "mode", mode, &len); if (e!=ESP_OK) strcpy(mode, "STA");
    uint8_t ch=36; nvs_get_u8(h, "mon_ch", &ch); *mon_ch = ch;
    uint8_t d=1; nvs_get_u8(h, "dhcp", &d); *dhcp = (d!=0);
    nvs_close(h);
    return true;
}

void wifi_cfg_force_dhcp(bool enable){ cfg_dhcp = enable; }

bool wifi_cfg_get_mode(char *mode, uint8_t *mon_ch) {
    if (!mode || !mon_ch) return false;
    nvs_handle_t h;
    if (nvs_open("netcfg", NVS_READONLY, &h) != ESP_OK) {
        strcpy(mode, "STA"); *mon_ch = 36; return false;
    }
    size_t len = 16;
    if (nvs_get_str(h, "mode", mode, &len) != ESP_OK) strcpy(mode, "STA");
    uint8_t ch = 36;
    nvs_get_u8(h, "mon_ch", &ch);
    *mon_ch = ch;
    nvs_close(h);
    return true;
}

static atomic_bool s_net_stack_ready = false;
static esp_err_t ensure_net_stack_once(void) {
    bool expected = false;
    if (atomic_compare_exchange_strong(&s_net_stack_ready, &expected, true)) {
        ESP_RETURN_ON_ERROR(esp_netif_init(), TAG, "esp_netif_init");
        esp_err_t err = esp_event_loop_create_default();
        if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
            ESP_RETURN_ON_ERROR(err, TAG, "event loop create");
        }
    }
    return ESP_OK;
}

static atomic_bool s_wifi_inited = false;
esp_err_t wifi_ensure_inited(void) {
    bool expected = false;
    if (!atomic_compare_exchange_strong(&s_wifi_inited, &expected, true)) return ESP_OK;
    ESP_RETURN_ON_ERROR(ensure_net_stack_once(), TAG, "net stack");
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    esp_err_t err = esp_wifi_init(&cfg);
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
        atomic_store(&s_wifi_inited, false);
        ESP_RETURN_ON_ERROR(err, TAG, "esp_wifi_init");
    }
    return ESP_OK;
}

static void apply_ip_static(const char* ip, const char* mask, const char* gw){
    if (!sta_netif) return;
    if (!ip || !ip[0] || !mask || !mask[0] || !gw || !gw[0]) return;
    esp_netif_ip_info_t info = {0};
    esp_netif_dhcpc_stop(sta_netif);
    info.ip.addr      = esp_ip4addr_aton(ip);
    info.netmask.addr = esp_ip4addr_aton(mask);
    info.gw.addr      = esp_ip4addr_aton(gw);
    ESP_ERROR_CHECK( esp_netif_set_ip_info(sta_netif, &info) );
}

bool wifi_cfg_apply_from_nvs(void) {
    char ssid[64]={0}, pass[64]={0}, ip[32]={0}, mask[32]={0}, gw[32]={0};
    char band[16]={0}, bw[16]={0}, powersave[16]={0}, mode[16]={0};
    uint8_t mon_ch = 36; bool dhcp = true;

    if (!load_cfg(ssid,sizeof(ssid), pass,sizeof(pass), ip,sizeof(ip), mask,sizeof(mask), gw,sizeof(gw),
                  band,sizeof(band), bw,sizeof(bw), powersave,sizeof(powersave), mode,sizeof(mode), &mon_ch, &dhcp)){
        return false;
    }
    if (ssid[0] == '\0') return false;

    static bool inited = false;
    if (!inited){
        nvs_flash_init();
        ensure_net_stack_once();
        if (sta_netif == NULL) sta_netif = esp_netif_create_default_wifi_sta();
        wifi_ensure_inited();
        inited = true;
    }

    wifi_config_t wcfg = {0};
    strlcpy((char*)wcfg.sta.ssid, ssid, sizeof(wcfg.sta.ssid));
    strlcpy((char*)wcfg.sta.password, pass, sizeof(wcfg.sta.password));
    wcfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wcfg.sta.sae_pwe_h2e = WPA3_SAE_PWE_BOTH;
    wcfg.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;

    if (strcmp(band, "5G") == 0) wcfg.sta.channel = 0;
    else wcfg.sta.channel = 0;

    esp_wifi_set_mode(WIFI_MODE_STA);

    #if CONFIG_IDF_TARGET_ESP32C5 || CONFIG_IDF_TARGET_ESP32C6
    wifi_protocols_t protocols = {
        .ghz_2g = WIFI_PROTOCOL_11B | WIFI_PROTOCOL_11G | WIFI_PROTOCOL_11N,
        .ghz_5g = WIFI_PROTOCOL_11A | WIFI_PROTOCOL_11N | WIFI_PROTOCOL_11AC | WIFI_PROTOCOL_11AX,
    };
    esp_wifi_set_protocols(WIFI_IF_STA, &protocols);
    #else
    esp_wifi_set_protocol(WIFI_IF_STA, WIFI_PROTOCOL_11B | WIFI_PROTOCOL_11G | WIFI_PROTOCOL_11N);
    #endif

    esp_wifi_set_config(WIFI_IF_STA, &wcfg);

    if (!dhcp && ip[0]) apply_ip_static(ip, mask, gw);
    else if (sta_netif) esp_netif_dhcpc_start(sta_netif);

    #if CONFIG_IDF_TARGET_ESP32C5 || CONFIG_IDF_TARGET_ESP32C6
    wifi_bandwidths_t bandwidths = {.ghz_2g = WIFI_BW_HT20, .ghz_5g = WIFI_BW_HT20};
    if (strcmp(bw, "VHT80") == 0) { bandwidths.ghz_2g = WIFI_BW_HT40; bandwidths.ghz_5g = WIFI_BW80; }
    else if (strcmp(bw, "HT40") == 0) { bandwidths.ghz_2g = WIFI_BW_HT40; bandwidths.ghz_5g = WIFI_BW_HT40; }
    esp_wifi_set_bandwidths(WIFI_IF_STA, &bandwidths);
    #else
    wifi_bandwidth_t bandwidth = WIFI_BW_HT20;
    if (strcmp(bw, "HT40") == 0) bandwidth = WIFI_BW_HT40;
    esp_wifi_set_bandwidth(WIFI_IF_STA, bandwidth);
    #endif

    esp_wifi_start();

    wifi_ps_type_t ps_mode = WIFI_PS_NONE;
    if (strcmp(powersave, "MIN") == 0) ps_mode = WIFI_PS_MIN_MODEM;
    else if (strcmp(powersave, "MAX") == 0) ps_mode = WIFI_PS_MAX_MODEM;
    esp_wifi_set_ps(ps_mode);

    esp_wifi_connect();
    return true;
}

// --- Parsing State ---
typedef struct {
    // Network
    char ssid[64];
    char pass[64];
    char ip[32];
    char mask[32];
    char gw[32];
    char band[16];
    char bw[16];
    char powersave[16];
    char mode[16];
    uint8_t mon_ch;
    bool dhcp;
    bool csi_enable;

    // Iperf
    char iperf_dest[32];
    char iperf_role[16];
    char iperf_proto[8];
    uint32_t iperf_period_us;
    uint32_t iperf_burst;
    uint32_t iperf_len;
    uint32_t iperf_port;
    bool iperf_enable;
} cfg_state_t;

static void on_cfg_line(const char *line, cfg_state_t *s){
    // Network Parsing
    if (strncmp(line, "SSID:",5)==0){ strncpy(s->ssid, line+5, 63); s->ssid[63]=0; return; }
    if (strncmp(line, "PASS:",5)==0){ strncpy(s->pass, line+5, 63); s->pass[63]=0; return; }
    if (strncmp(line, "IP:",3)==0){   strncpy(s->ip,   line+3, 31); s->ip[31]=0;   return; }
    if (strncmp(line, "MASK:",5)==0){ strncpy(s->mask, line+5, 31); s->mask[31]=0; return; }
    if (strncmp(line, "GW:",3)==0){   strncpy(s->gw,   line+3, 31); s->gw[31]=0;   return; }
    if (strncmp(line, "BAND:",5)==0){ strncpy(s->band, line+5, 15); s->band[15]=0; return; }
    if (strncmp(line, "BW:",3)==0){   strncpy(s->bw,   line+3, 15); s->bw[15]=0;   return; }
    if (strncmp(line, "POWERSAVE:",10)==0){ strncpy(s->powersave, line+10, 15); s->powersave[15]=0; return; }
    if (strncmp(line, "MODE:",5)==0){ strncpy(s->mode, line+5, 15); s->mode[15]=0; return; }
    if (strncmp(line, "MON_CH:",7)==0){ s->mon_ch = atoi(line+7); return; }
    if (strncmp(line, "DHCP:",5)==0){ s->dhcp = atoi(line+5) ? true:false; return; }

    // Iperf Parsing (Matches Python Script Keys)
    // Support both DEST and DST to be safe
    if (strncmp(line, "IPERF_DEST_IP:", 14) == 0) { strncpy(s->iperf_dest, line+14, 31); s->iperf_dest[31]=0; return; }
    if (strncmp(line, "IPERF_DST_IP:", 13) == 0)  { strncpy(s->iperf_dest, line+13, 31); s->iperf_dest[31]=0; return; }

    if (strncmp(line, "IPERF_ROLE:", 11) == 0)    { strncpy(s->iperf_role, line+11, 15); s->iperf_role[15]=0; return; }
    if (strncmp(line, "IPERF_PROTO:", 12) == 0)   { strncpy(s->iperf_proto, line+12, 7); s->iperf_proto[7]=0; return; }

    if (strncmp(line, "IPERF_PERIOD_US:", 16) == 0) { s->iperf_period_us = atoi(line+16); return; }
    if (strncmp(line, "IPERF_BURST:", 12) == 0)     { s->iperf_burst = atoi(line+12); return; }
    if (strncmp(line, "IPERF_LEN:", 10) == 0)       { s->iperf_len = atoi(line+10); return; }
    if (strncmp(line, "IPERF_PORT:", 11) == 0)      { s->iperf_port = atoi(line+11); return; }
    if (strncmp(line, "IPERF_ENABLED:", 14) == 0)   { s->iperf_enable = atoi(line+14) ? true:false; return; }

#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
    if (strncmp(line, "CSI:",4)==0){ s->csi_enable = atoi(line+4) ? true:false; return; }
#endif
}

static bool wifi_cfg_cmd_handler(const char *line, cmd_reply_func_t reply_func, void *reply_ctx) {
    static bool in_cfg = false;
    static cfg_state_t s;

    if (!in_cfg) {
        if (strcmp(line, "CFG") == 0) {
            in_cfg = true;
            // Clear all buffers
            memset(&s, 0, sizeof(s));
            // Set Defaults
            s.mon_ch = 36;
            s.dhcp = true;
            s.csi_enable = false;
            s.iperf_period_us = 10000;
            s.iperf_burst = 1;
            s.iperf_len = 1470;
            s.iperf_port = 5001;
            s.iperf_enable = true;
            strcpy(s.iperf_dest, "192.168.1.50");
            strcpy(s.iperf_role, "CLIENT");
            strcpy(s.iperf_proto, "UDP");

            // --- NEW: ACK Start of CFG ---
            if (reply_func) reply_func("OK\n", reply_ctx);
            return true;
        }
        return false;
    }

    if (strcmp(line, "END") == 0) {
        // Apply Network Defaults if missing
        if (!s.band[0]) strcpy(s.band, "2.4G");
        if (!s.bw[0])   strcpy(s.bw, "HT20");
        if (!s.powersave[0]) strcpy(s.powersave, "NONE");
        if (!s.mode[0]) strcpy(s.mode, "STA");

        // 1. Save Network
        save_net_cfg(s.ssid, s.pass, s.ip, s.mask, s.gw, s.dhcp, s.band, s.bw, s.powersave, s.mode, s.mon_ch);

        // 2. Save Iperf
        save_iperf_cfg(s.iperf_dest, s.iperf_role, s.iperf_proto, s.iperf_period_us, s.iperf_burst, s.iperf_len, s.iperf_port, s.iperf_enable);

        // 3. Save CSI
#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
        csi_mgr_save_enable_state(s.csi_enable);
#endif

        // --- NEW: ACK End of CFG ---
        if (reply_func) reply_func("OK Config Saved\n", reply_ctx);

        // Apply changes immediately
        wifi_cfg_apply_from_nvs();

        in_cfg = false;
        return true;
    }

    on_cfg_line(line, &s);

    // --- NEW: ACK Intermediate Line ---
    if (reply_func) reply_func("OK\n", reply_ctx);

    return true;
}

void wifi_cfg_init(void){
    nvs_flash_init();
    cmd_transport_init();
    cmd_transport_register_listener(wifi_cfg_cmd_handler);
}

wifi_ps_type_t wifi_cfg_get_power_save_mode(void) {
    char powersave[16] = {0};
    nvs_handle_t h;
    if (nvs_open("netcfg", NVS_READONLY, &h) != ESP_OK) return WIFI_PS_NONE;
    size_t len = sizeof(powersave);
    nvs_get_str(h, "powersave", powersave, &len);
    nvs_close(h);

    if (strcmp(powersave, "MIN") == 0) return WIFI_PS_MIN_MODEM;
    if (strcmp(powersave, "MAX") == 0) return WIFI_PS_MAX_MODEM;
    return WIFI_PS_NONE;
}

bool wifi_cfg_get_bandwidth(char *buf, size_t buf_size) {
    if (!buf || buf_size < 1) return false;
    nvs_handle_t h;
    if (nvs_open("netcfg", NVS_READONLY, &h) != ESP_OK) {
        strncpy(buf, "Unknown", buf_size); return false;
    }
    size_t len = buf_size;
    esp_err_t err = nvs_get_str(h, "bw", buf, &len);
    nvs_close(h);
    return (err == ESP_OK);
}