#include "cmd_transport.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";

#define MAX_LISTENERS 4
static cmd_line_handler_t s_listeners[MAX_LISTENERS] = {0};
static int s_listener_count = 0;
static bool s_inited = false;

void cmd_transport_register_listener(cmd_line_handler_t handler) {
    if (s_listener_count < MAX_LISTENERS) {
        s_listeners[s_listener_count++] = handler;
    }
}

static void trim_trailing(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;
    }
}

// --- UPDATED DISPATCH LOGIC ---
static void dispatch_line(char *line, cmd_reply_func_t reply_func, void *reply_ctx) {
    bool handled = false;

    // 1. Offer to registered listeners (e.g. wifi_cfg)
    // Note: The listener (wifi_cfg) is now responsible for sending "OK" if it handles the line.
    for (int i = 0; i < s_listener_count; i++) {
        if (s_listeners[i] && s_listeners[i](line, reply_func, reply_ctx)) {
            handled = true;
            break;
        }
    }

    // 2. If not handled, pass to system console
    if (!handled && strlen(line) > 0) {
        int ret;
        esp_err_t err = esp_console_run(line, &ret);

        if (err == ESP_ERR_NOT_FOUND) {
            // Robustness: Always reply, even for unknown commands
            if (reply_func) reply_func("ERROR: Unknown Command\n", reply_ctx);
        } else if (err == ESP_OK) {
            if (ret == 0) {
                // Command Success -> Send ACK so Python knows to proceed
                if (reply_func) reply_func("OK\n", reply_ctx);
            } else {
                // Command logic failed (e.g., iperf returned 1)
                char buf[64];
                snprintf(buf, sizeof(buf), "ERROR: Command failed (ret=%d)\n", ret);
                if (reply_func) reply_func(buf, reply_ctx);
            }
        } else {
            if (reply_func) reply_func("ERROR: System execution failed\n", reply_ctx);
        }
    }
}

// --- UART (stdin/stdout) Support ---
static void uart_reply(const char *msg, void *ctx) {
    (void)ctx;
    printf("%s", msg);
    fflush(stdout);
}

static void uart_listener_task(void *arg) {
    char line[256];
    setvbuf(stdin, NULL, _IONBF, 0);
    setvbuf(stdout, NULL, _IONBF, 0);

    while (1) {
        if (fgets(line, sizeof(line), stdin)) {
            trim_trailing(line);
            dispatch_line(line, uart_reply, NULL);
        } else {
            vTaskDelay(pdMS_TO_TICKS(20));
        }
    }
}

// --- USB Serial/JTAG Support ---
#if SOC_USB_SERIAL_JTAG_SUPPORTED
static void usb_reply(const char *msg, void *ctx) {
    (void)ctx;
    usb_serial_jtag_write_bytes((const uint8_t*)msg, strlen(msg), pdMS_TO_TICKS(50));
}

static void usb_listener_task(void *arg) {
    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-Serial/JTAG driver");
        vTaskDelete(NULL);
    }

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

    while (1) {
        int n = usb_serial_jtag_read_bytes(&c, 1, pdMS_TO_TICKS(20));
        if (n > 0) {
            if (c == '\n' || c == '\r') {
                if (idx > 0) {
                    buf[idx] = 0;
                    dispatch_line(buf, usb_reply, 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;
    xTaskCreatePinnedToCore(uart_listener_task, "cmd_uart", 4096, NULL, 5, NULL, tskNO_AFFINITY);
    #if SOC_USB_SERIAL_JTAG_SUPPORTED
    xTaskCreatePinnedToCore(usb_listener_task, "cmd_usb", 4096, NULL, 5, NULL, tskNO_AFFINITY);
    #endif
}