// wifi_cfg_dualmode.c — ESP-IDF v5.3.x
// Listens for CFG/END Wi‑Fi config on BOTH UART(console/COM) and USB‑Serial/JTAG concurrently.
//  - UART path uses stdio (fgets on stdin) — works with /dev/ttyUSB*
//  - USB path uses driver API (usb_serial_jtag_read_bytes / write_bytes) — works with /dev/ttyACM*
//  - Tolerates ESP_ERR_INVALID_STATE on repeated inits
//  - Supports DHCP or static IP, persists to NVS, applies immediately
//  - Bandwidth options: HT20 (20MHz), HT40 (40MHz), VHT80 (80MHz, 5GHz only)
//  - Power save modes: NONE (default, best for CSI), MIN/MIN_MODEM, MAX/MAX_MODEM

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

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.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 "esp_event.h"
#include "driver/usb_serial_jtag.h"   // direct read/write API (no VFS remap)
#include "wifi_cfg.h"


// wifi_cfg.c

static const char *TAG = "wifi_cfg";

static esp_netif_t *sta_netif = NULL;
static bool cfg_dhcp = true;

static void trim(char *s){
    int n = strlen(s);
    while(n>0 && (s[n-1]=='\r' || s[n-1]=='\n' || isspace((unsigned char)s[n-1]))) s[--n]=0;
    while(*s && isspace((unsigned char)*s)){
        memmove(s, s+1, strlen(s));
    }
}

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

static void save_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){
    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);
    nvs_set_u8(h, "dhcp", dhcp ? 1 : 0);
    nvs_commit(h);
    nvs_close(h);
    cfg_dhcp = dhcp;
    ESP_LOGI(TAG, "Config saved to NVS: SSID=%s Band=%s BW=%s PowerSave=%s", ssid?ssid:"", band?band:"", bw?bw:"", powersave?powersave:"NONE");
}

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, 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");
    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; }


/* --- One-time net stack bring-up (thread-safe) --- */
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;
}

/* --- One-time Wi-Fi driver init (thread-safe, idempotent) --- */
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)) {
        // someone else already initialized (or is initializing and finished)
        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) {
        // roll back the flag so a later attempt can retry
        atomic_store(&s_wifi_inited, false);
        ESP_RETURN_ON_ERROR(err, TAG, "esp_wifi_init");
    }

    // Optional: set default interface and mode now or let caller do it
    // ESP_RETURN_ON_ERROR(esp_wifi_set_mode(WIFI_MODE_STA), TAG, "set mode");

    return ESP_OK;
}


static void apply_ip_static(const char* ip, const char* mask, const char* gw){
    if (!sta_netif) return;

    // Validate IP strings are not empty
    if (!ip || !ip[0] || !mask || !mask[0] || !gw || !gw[0]) {
        ESP_LOGW(TAG, "Invalid static IP config: IP=%s MASK=%s GW=%s",
                 ip ? ip : "NULL", mask ? mask : "NULL", gw ? gw : "NULL");
        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_LOGI(TAG, "Setting static IP: %s, netmask: %s, gateway: %s", ip, mask, 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};
    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), &dhcp)){
        ESP_LOGW(TAG, "No Wi‑Fi config in NVS");
        return false;
    }
    // Treat empty SSID as "no config" to avoid ESP_ERR_WIFI_SSID panics
    if (ssid[0] == '\0') {
        ESP_LOGW(TAG, "SSID in NVS is empty; treating as no Wi-Fi config");
        return false;
    }
    ESP_LOGI(TAG, "Applying Wi-Fi config: SSID=%s DHCP=%d IP=%s MASK=%s GW=%s Band=%s BW=%s PowerSave=%s",
             ssid, dhcp, ip, mask, gw, band, bw, powersave);

    static bool inited = false;
    if (!inited){
        // NVS (with recovery)
        esp_err_t err = nvs_flash_init();
        if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
            nvs_flash_erase();
            err = nvs_flash_init();
        }
        if (err != ESP_OK) { ESP_ERROR_CHECK(err); }

        // Netif + default event loop (tolerate already-initialized state)
        do{ esp_err_t __e = esp_netif_init(); if(__e!=ESP_OK && __e!=ESP_ERR_INVALID_STATE){ ESP_ERROR_CHECK(__e);} }while(0);
        if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) { ESP_ERROR_CHECK(err); }

        do{ esp_err_t __e = esp_event_loop_create_default(); if(__e!=ESP_OK && __e!=ESP_ERR_INVALID_STATE){ ESP_ERROR_CHECK(__e);} }while(0);
        if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) { ESP_ERROR_CHECK(err); }

        if (sta_netif == NULL) {
            sta_netif = esp_netif_create_default_wifi_sta();
        }

        ESP_ERROR_CHECK(err=wifi_ensure_inited());
        if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) { ESP_ERROR_CHECK(err); }

        inited = true;
    }

    wifi_config_t wcfg = (wifi_config_t){0};
    strncpy((char*)wcfg.sta.ssid, ssid, sizeof(wcfg.sta.ssid)-1);
    strncpy((char*)wcfg.sta.password, pass, sizeof(wcfg.sta.password)-1);
    wcfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wcfg.sta.sae_pwe_h2e = WPA3_SAE_PWE_BOTH;

    // Set scan method to search all channels
    wcfg.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;

    // Log and configure band preference
    if (strcmp(band, "5G") == 0) {
        ESP_LOGI(TAG, "Configuring for 5GHz operation");
        // For 5GHz preference, scan both but prefer 5GHz channels
        wcfg.sta.channel = 0;  // Scan all channels (both 2.4GHz and 5GHz)
    } else {
        ESP_LOGI(TAG, "Configuring for 2.4GHz operation (default)");
        wcfg.sta.channel = 0;  // Scan all channels
    }

    ESP_ERROR_CHECK( esp_wifi_set_mode(WIFI_MODE_STA) );

    // Configure WiFi protocols based on chip capabilities
#if CONFIG_IDF_TARGET_ESP32C5 || CONFIG_IDF_TARGET_ESP32C6
    // Dual-band chips (2.4GHz + 5GHz support)
    wifi_protocols_t protocols = {
        // 2.4 GHz: b/g/n (no ax for CSI compatibility)
        .ghz_2g = WIFI_PROTOCOL_11B |
                  WIFI_PROTOCOL_11G |
                  WIFI_PROTOCOL_11N,
        // 5 GHz: a/n/ac/ax
        .ghz_5g = WIFI_PROTOCOL_11A |
                  WIFI_PROTOCOL_11N |
                  WIFI_PROTOCOL_11AC |
                  WIFI_PROTOCOL_11AX,
    };
    ESP_ERROR_CHECK( esp_wifi_set_protocols(WIFI_IF_STA, &protocols) );
#else
    // Single-band chips (2.4GHz only) - ESP32, ESP32-S2, ESP32-S3, ESP32-C3, etc.
    // Use legacy API for 2.4GHz-only chips
    uint8_t protocol_bitmap = WIFI_PROTOCOL_11B | WIFI_PROTOCOL_11G | WIFI_PROTOCOL_11N;
    ESP_ERROR_CHECK( esp_wifi_set_protocol(WIFI_IF_STA, protocol_bitmap) );

    // Warn if user requested 5GHz on a 2.4GHz-only chip
    if (strcmp(band, "5G") == 0) {
        ESP_LOGW(TAG, "5GHz requested but this chip only supports 2.4GHz - using 2.4GHz");
    }
#endif

    ESP_ERROR_CHECK( esp_wifi_set_config(WIFI_IF_STA, &wcfg) );

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

    // Set bandwidth BEFORE WiFi is started
#if CONFIG_IDF_TARGET_ESP32C5 || CONFIG_IDF_TARGET_ESP32C6
    // Dual-band chips use struct-based API
    wifi_bandwidths_t bandwidths = {
        .ghz_2g = WIFI_BW_HT20,
        .ghz_5g = WIFI_BW_HT20
    };

    if (strcmp(bw, "VHT80") == 0) {
        // 80MHz only supported on 5GHz
        bandwidths.ghz_2g = WIFI_BW_HT40;  // 2.4GHz fallback to 40MHz
        bandwidths.ghz_5g = WIFI_BW80;     // Use WIFI_BW80, not WIFI_BW_HT80
        ESP_LOGI(TAG, "Setting bandwidth to VHT80 (80MHz) for 5GHz, HT40 (40MHz) for 2.4GHz");
    } else if (strcmp(bw, "HT40") == 0) {
        bandwidths.ghz_2g = WIFI_BW_HT40;
        bandwidths.ghz_5g = WIFI_BW_HT40;
        ESP_LOGI(TAG, "Setting bandwidth to HT40 (40MHz) for both bands");
    } else {
        ESP_LOGI(TAG, "Setting bandwidth to HT20 (20MHz) for both bands");
    }

    esp_err_t bw_err = esp_wifi_set_bandwidths(WIFI_IF_STA, &bandwidths);
    if (bw_err != ESP_OK) {
        ESP_LOGW(TAG, "Failed to set bandwidths: %s", esp_err_to_name(bw_err));
    }
#else
    // Single-band chips (2.4GHz only) use legacy API
    wifi_bandwidth_t bandwidth = WIFI_BW_HT20;

    if (strcmp(bw, "VHT80") == 0) {
        ESP_LOGW(TAG, "VHT80 (80MHz) not supported on 2.4GHz-only chips - using HT40 (40MHz)");
        bandwidth = WIFI_BW_HT40;
    } else if (strcmp(bw, "HT40") == 0) {
        bandwidth = WIFI_BW_HT40;
        ESP_LOGI(TAG, "Setting bandwidth to HT40 (40MHz) for 2.4GHz");
    } else {
        ESP_LOGI(TAG, "Setting bandwidth to HT20 (20MHz) for 2.4GHz");
    }

    esp_err_t bw_err = esp_wifi_set_bandwidth(WIFI_IF_STA, bandwidth);
    if (bw_err != ESP_OK) {
        ESP_LOGW(TAG, "Failed to set bandwidth: %s", esp_err_to_name(bw_err));
    }
#endif

    esp_err_t err2 = esp_wifi_start();
    if (err2 != ESP_OK && err2 != ESP_ERR_INVALID_STATE) { ESP_ERROR_CHECK(err2); }

    // Set power save mode based on configuration
    wifi_ps_type_t ps_mode = WIFI_PS_NONE;  // Default: no power save (best for CSI)

    if (strcmp(powersave, "MIN") == 0 || strcmp(powersave, "MIN_MODEM") == 0) {
        ps_mode = WIFI_PS_MIN_MODEM;
        ESP_LOGI(TAG, "Setting power save mode: MIN_MODEM");
    } else if (strcmp(powersave, "MAX") == 0 || strcmp(powersave, "MAX_MODEM") == 0) {
        ps_mode = WIFI_PS_MAX_MODEM;
        ESP_LOGI(TAG, "Setting power save mode: MAX_MODEM");
    } else {
        ESP_LOGI(TAG, "Setting power save mode: NONE (best for CSI)");
    }

    err2 = esp_wifi_set_ps(ps_mode);
    if (err2 != ESP_OK) {
        ESP_LOGW(TAG, "Failed to set power save mode: %s", esp_err_to_name(err2));
    }

    err2 = esp_wifi_connect();
    if (err2 != ESP_OK && err2 != ESP_ERR_WIFI_NOT_INIT && err2 != ESP_ERR_INVALID_STATE) { ESP_ERROR_CHECK(err2); }
    return true;
}

// -------------------- Dual input paths --------------------

typedef struct {
    // emit() should write a short line back to the same channel (optional)
    void (*emit)(const char *s, void *ctx);
    void *ctx;
    // fetch_line() should fill buf with a single line (without trailing CR/LF) and return true if a line was read
    bool (*fetch_line)(char *buf, size_t sz, void *ctx);
} cfg_io_t;

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

static void cfg_worker(const cfg_io_t *io){
    char line[160];
    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};
    bool dhcp = true;
    bool in_cfg = false;

    for(;;){
        if (!io->fetch_line(line, sizeof(line), io->ctx)){
            vTaskDelay(pdMS_TO_TICKS(20));
            continue;
        }
        trim(line);
        if (!in_cfg){
            if (strcmp(line, "CFG")==0){
                in_cfg = true;
                ssid[0]=pass[0]=ip[0]=mask[0]=gw[0]=0;
                band[0]=bw[0]=powersave[0]=0;
                dhcp = true;
            }
            continue;
        }
        if (strcmp(line, "END")==0){
            // Set defaults if not specified
            if (!band[0]) strcpy(band, "2.4G");
            if (!bw[0]) strcpy(bw, "HT20");
            if (!powersave[0]) strcpy(powersave, "NONE");

            save_cfg(ssid, pass, ip, mask, gw, dhcp, band, bw, powersave);
            if (io->emit) io->emit("OK\n", io->ctx);
            wifi_cfg_apply_from_nvs();
            in_cfg = false;
            continue;
        }
        on_cfg_line(line, ssid, pass, ip, mask, gw, band, bw, powersave, &dhcp);
    }
}

// ---- UART(stdin) path ----
static bool uart_fetch_line(char *buf, size_t sz, void *ctx){
    (void)ctx;
    if (!fgets(buf, sz, stdin)) return false;
    return true;
}
static void uart_emit(const char *s, void *ctx){
    (void)ctx;
    fputs(s, stdout);
    fflush(stdout);
}
static void cfg_listener_uart_task(void *arg){
    setvbuf(stdin,  NULL, _IONBF, 0);
    setvbuf(stdout, NULL, _IONBF, 0);
    cfg_io_t io = {.emit = uart_emit, .ctx = NULL, .fetch_line = uart_fetch_line};
    cfg_worker(&io);
}

// ---- USB-Serial/JTAG path ----
typedef struct { int dummy; } usb_ctx_t;
static bool usb_fetch_line(char *buf, size_t sz, void *ctx){
    (void)ctx;
    static char acc[256];
    static size_t acc_len = 0;

    uint8_t tmp[64];
    int n = usb_serial_jtag_read_bytes(tmp, sizeof(tmp), pdMS_TO_TICKS(10));
    if (n <= 0){
        return false;
    }
    for (int i=0; i<n; ++i){
        char c = (char)tmp[i];
        if (c == '\n' || c == '\r'){
            acc[acc_len] = 0;
            strncpy(buf, acc, sz-1);
            buf[sz-1]=0;
            acc_len = 0;
            return true;
        }else if (acc_len < sizeof(acc)-1){
            acc[acc_len++] = c;
        }
    }
    return false;
}
static void usb_emit(const char *s, void *ctx){
    (void)ctx;
    usb_serial_jtag_write_bytes((const uint8_t*)s, strlen(s), pdMS_TO_TICKS(50));
}
static void cfg_listener_usb_task(void *arg){
    // Install USB-Serial/JTAG driver (no VFS remap: stdio stays on UART)
    usb_serial_jtag_driver_config_t d = USB_SERIAL_JTAG_DRIVER_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(usb_serial_jtag_driver_install(&d));

    usb_ctx_t uctx = {0};
    cfg_io_t io = {.emit = usb_emit, .ctx = &uctx, .fetch_line = usb_fetch_line};
    cfg_worker(&io);
}

void wifi_cfg_init(void){
    // Make sure NVS exists
    nvs_flash_init();

    // Spawn both listeners
    xTaskCreatePinnedToCore(cfg_listener_uart_task, "cfg_uart",
                            6144, NULL, 5, NULL, tskNO_AFFINITY);
    xTaskCreatePinnedToCore(cfg_listener_usb_task, "cfg_usb",
                            6144, NULL, 5, NULL, tskNO_AFFINITY);
}

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) {
        ESP_LOGW(TAG, "Failed to open NVS for power save - defaulting to NONE");
        return WIFI_PS_NONE;
    }

    size_t len = sizeof(powersave);
    esp_err_t err = nvs_get_str(h, "powersave", powersave, &len);
    nvs_close(h);

    if (err != ESP_OK) {
        ESP_LOGI(TAG, "No power save setting in NVS - defaulting to NONE");
        return WIFI_PS_NONE;
    }

    // Parse the power save string
    if (strcmp(powersave, "MIN") == 0 || strcmp(powersave, "MIN_MODEM") == 0) {
        ESP_LOGI(TAG, "Power save mode: MIN_MODEM");
        return WIFI_PS_MIN_MODEM;
    } else if (strcmp(powersave, "MAX") == 0 || strcmp(powersave, "MAX_MODEM") == 0) {
        ESP_LOGI(TAG, "Power save mode: MAX_MODEM");
        return WIFI_PS_MAX_MODEM;
    } else {
        ESP_LOGI(TAG, "Power save mode: NONE (no power saving)");
        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 - 1);
        buf[buf_size - 1] = '\0';
        return false;
    }

    size_t len = buf_size;
    esp_err_t err = nvs_get_str(h, "bw", buf, &len);
    nvs_close(h);

    if (err != ESP_OK) {
        strncpy(buf, "HT20", buf_size - 1);  // Default
        buf[buf_size - 1] = '\0';
        return false;
    }

    return true;
}