// wifi_cfg.c - Refactored to use cmd_transport
#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" // Dependency

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

// --- NVS & Helper Functions (Kept from original) ---

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, 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, "Config saved to NVS: SSID=%s Mode=%s MonCh=%d", ssid?ssid:"", mode?mode:"STA", mon_ch);
}

// ... [Existing load_cfg, apply_ip_static, wifi_ensure_inited, wifi_cfg_apply_from_nvs implementations retained here] ...
// (Omitting full copy-paste of unchanged helper functions for brevity, assume they are present as in your original upload)

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

// ... [Include apply_ip_static, ensure_net_stack_once, wifi_ensure_inited, wifi_cfg_apply_from_nvs from original file] ...
// (You must keep these functions in the final file)

// --- Logic for parsing the config lines ---

static void on_cfg_line(const char *line, char *ssid, char *pass, char *ip, char *mask, char *gw, char *band, char *bw, char *powersave, char *mode, uint8_t *mon_ch, 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, "MODE:",5)==0){ strncpy(mode, line+5, 15); mode[15]=0; return; }
    if (strncmp(line, "MON_CH:",7)==0){ *mon_ch = atoi(line+7); return; }
    if (strncmp(line, "DHCP:",5)==0){ *dhcp = atoi(line+5) ? true:false;      return; }
}

// --- The Listener Callback (Registered with cmd_transport) ---

static bool wifi_cfg_line_handler(const char *line, cmd_transport_reply_t reply_func, void *reply_ctx) {
    // State is static because we receive one line at a time from the transport task
    static bool in_cfg = false;
    static char ssid[64]={0}, pass[64]={0}, ip[32]={0}, mask[32]={0}, gw[32]={0};
    static char band[16]={0}, bw[16]={0}, powersave[16]={0}, mode[16]={0};
    static uint8_t mon_ch = 36;
    static bool dhcp = true;

    // Check for start of config block
    if (strcmp(line, "CFG") == 0) {
        in_cfg = true;
        // Reset buffers
        ssid[0]=0; pass[0]=0; ip[0]=0; mask[0]=0; gw[0]=0;
        band[0]=0; bw[0]=0; powersave[0]=0; mode[0]=0;
        mon_ch = 36; dhcp = true;
        return true; // We handled this line
    }

    if (in_cfg) {
        if (strcmp(line, "END") == 0) {
            // Apply defaults
            if (!band[0]) strcpy(band, "2.4G");
            if (!bw[0]) strcpy(bw, "HT20");
            if (!powersave[0]) strcpy(powersave, "NONE");
            if (!mode[0]) strcpy(mode, "STA");

            save_cfg(ssid, pass, ip, mask, gw, dhcp, band, bw, powersave, mode, mon_ch);
            
            // Reply "OK" to whoever sent the commands
            if (reply_func) {
                reply_func("OK\n", reply_ctx);
            }
            
            // Trigger application of new config
            // Note: This might block briefly, ideally offload to event loop if heavy
            // For now we assume apply_from_nvs is fast enough or acceptable
            // To be safe, we could signal an event, but let's keep logic preserved.
            wifi_cfg_apply_from_nvs();
            
            in_cfg = false;
            return true;
        }
        
        // Parse config line
        on_cfg_line(line, ssid, pass, ip, mask, gw, band, bw, powersave, mode, &mon_ch, &dhcp);
        return true; // Consumed
    }

    return false; // Not in config mode, let other listeners (console) handle it
}

// --- Init ---

void wifi_cfg_init(void){
    // Ensure NVS is ready
    nvs_flash_init();

    // Initialize the transport layer (starts UART/USB tasks)
    cmd_transport_init();

    // Register our handler
    cmd_transport_register_listener(wifi_cfg_line_handler);
}

// ... [Keep wifi_cfg_get_power_save_mode, wifi_cfg_get_bandwidth, wifi_cfg_get_mode, wifi_cfg_force_dhcp, wifi_ensure_inited] ...

#include "cmd_transport.h"
#include "esp_system.h"
#include "esp_log.h"
#include "esp_console.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "soc/soc_caps.h"
#include <stdio.h>
#include <string.h>
#include <ctype.h>

#if SOC_USB_SERIAL_JTAG_SUPPORTED
    #include "driver/usb_serial_jtag.h"
#endif

static const char *TAG = "CMD_TP";

// Max listeners
#define MAX_LISTENERS 4
static cmd_line_callback_t s_listeners[MAX_LISTENERS] = {0};
static int s_listener_count = 0;
static bool s_inited = false;

// Register a listener
void cmd_transport_register_listener(cmd_line_callback_t cb) {
    if (s_listener_count < MAX_LISTENERS) {
        s_listeners[s_listener_count++] = cb;
    } else {
        ESP_LOGE(TAG, "Max listeners reached");
    }
}

// Helper: Trim whitespace
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;
    // Note: We don't trim leading whitespace to preserve command indentation if needed,
    // but the original code did. We'll stick to trailing trim for safety.
}

// Dispatch logic
static void dispatch_line(char *line, cmd_transport_reply_t reply_func, void *reply_ctx) {
    bool handled = false;

    // 1. Offer to registered listeners (e.g., wifi_cfg)
    for (int i = 0; i < s_listener_count; i++) {
        if (s_listeners[i] && s_listeners[i](line, reply_func, reply_ctx)) {
            handled = true;
            break; // Consumed by a listener
        }
    }

    // 2. If not consumed, pass to ESP Console (for 'mode_monitor', etc.)
    if (!handled && strlen(line) > 0) {
        int ret;
        esp_err_t err = esp_console_run(line, &ret);
        if (err == ESP_ERR_NOT_FOUND) {
            // Unrecognized command
            // Optional: reply_func("Unrecognized command\n", reply_ctx);
        } else if (err != ESP_OK) {
            ESP_LOGE(TAG, "Console run error: %d", err);
        }
    }
}

// ---- UART (stdin/stdout) Support ----
static void uart_emit(const char *s, void *ctx) {
    (void)ctx;
    fputs(s, stdout);
    fflush(stdout);
}

static void uart_listener_task(void *arg) {
    char line[256];
    // Disable buffering for immediate interaction
    setvbuf(stdin, NULL, _IONBF, 0);
    setvbuf(stdout, NULL, _IONBF, 0);

    while (1) {
        if (fgets(line, sizeof(line), stdin)) {
            trim(line);
            dispatch_line(line, uart_emit, NULL);
        } else {
            vTaskDelay(pdMS_TO_TICKS(50));
        }
    }
}

// ---- USB Serial/JTAG Support ----
#if SOC_USB_SERIAL_JTAG_SUPPORTED
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 usb_listener_task(void *arg) {
    // Init driver
    usb_serial_jtag_driver_config_t d = USB_SERIAL_JTAG_DRIVER_CONFIG_DEFAULT();
    if (usb_serial_jtag_driver_install(&d) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to install USB driver");
        vTaskDelete(NULL);
    }

    char buf[256];
    size_t idx = 0;
    uint8_t c;

    while (1) {
        // Read byte-by-byte to construct lines
        int n = usb_serial_jtag_read_bytes(&c, 1, pdMS_TO_TICKS(50));
        if (n > 0) {
            if (c == '\n' || c == '\r') {
                if (idx > 0) {
                    buf[idx] = 0;
                    dispatch_line(buf, usb_emit, NULL);
                    idx = 0;
                }
            } else {
                if (idx < sizeof(buf) - 1) {
                    buf[idx++] = (char)c;
                }
            }
        }
    }
}
#endif

void cmd_transport_init(void) {
    if (s_inited) return;
    s_inited = true;

    // Start UART listener (Always available)
    xTaskCreatePinnedToCore(uart_listener_task, "cmd_uart", 4096, NULL, 5, NULL, tskNO_AFFINITY);

    // Start USB listener (If supported)
    #if SOC_USB_SERIAL_JTAG_SUPPORTED
    xTaskCreatePinnedToCore(usb_listener_task, "cmd_usb", 4096, NULL, 5, NULL, tskNO_AFFINITY);
    #endif
}