#include "app_console.h"
#include "esp_console.h"
#include "esp_log.h"
#include "argtable3/argtable3.h"
#include "wifi_cfg.h"
#include "iperf.h"
#include "wifi_controller.h" // Required for Mode Switching
#include <string.h>
#include <arpa/inet.h>

// Helper to refresh prompt at end of commands
static void end_cmd(void) {
    app_console_update_prompt();
}

// ============================================================================
// COMMAND: iperf
// ============================================================================
static struct {
    struct arg_lit *start, *stop, *status, *save, *reload;
    struct arg_str *ip;
    struct arg_int *port, *pps, *len, *burst;
    struct arg_lit *help;
    struct arg_end *end;
} iperf_args;

static int cmd_iperf(int argc, char **argv) {
    int nerrors = arg_parse(argc, argv, (void **)&iperf_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, iperf_args.end, argv[0]);
        return 1;
    }

    if (iperf_args.help->count > 0) {
        printf("Usage: iperf [options]\n");
        printf("  --start       Start traffic\n");
        printf("  --stop        Stop traffic\n");
        printf("  --status      Show stats\n");
        printf("  --save        Save config to NVS\n");
        printf("  -c <ip>       Set Destination IP\n");
        printf("  --pps <n>     Set Packets Per Second\n");
        return 0;
    }

    if (iperf_args.reload->count > 0) {
        iperf_param_init();
        printf("Configuration reloaded from NVS.\n");
    }

    bool config_changed = false;
    iperf_cfg_t cfg;
    iperf_param_get(&cfg);

    if (iperf_args.ip->count > 0)    { cfg.dip = inet_addr(iperf_args.ip->sval[0]); config_changed = true; }
    if (iperf_args.port->count > 0)  { cfg.dport = (uint16_t)iperf_args.port->ival[0]; config_changed = true; }
    if (iperf_args.len->count > 0)   { cfg.send_len = (uint32_t)iperf_args.len->ival[0]; config_changed = true; }
    if (iperf_args.burst->count > 0) { cfg.burst_count = (uint32_t)iperf_args.burst->ival[0]; config_changed = true; }
    if (iperf_args.pps->count > 0)   {
        int pps = iperf_args.pps->ival[0];
        if (pps > 0) {
            cfg.target_pps = (uint32_t)pps;
            config_changed = true;
        }
    }

    if (config_changed) {
        iperf_param_set(&cfg);
        printf("RAM configuration updated.\n");
    }

    if (iperf_args.save->count > 0) {
        bool changed = false;
        if (iperf_param_save(&changed) == ESP_OK) {
            printf(changed ? "Configuration saved to NVS.\n" : "No changes to save (NVS matches RAM).\n");
        } else {
            printf("Error saving to NVS.\n");
        }
    }

    if (iperf_args.stop->count > 0) iperf_stop();
    if (iperf_args.start->count > 0) iperf_start();
    if (iperf_args.status->count > 0) iperf_print_status();

    end_cmd();
    return 0;
}

// ============================================================================
// COMMAND: monitor (Switch STA/MONITOR)
// ============================================================================
static struct {
    struct arg_lit *start;
    struct arg_lit *stop;
    struct arg_lit *status;
    struct arg_int *channel;
    struct arg_lit *help;
    struct arg_end *end;
} mon_args;

static int cmd_monitor(int argc, char **argv) {
    int nerrors = arg_parse(argc, argv, (void **)&mon_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, mon_args.end, argv[0]);
        return 1;
    }

    if (mon_args.help->count > 0) {
        printf("Usage: monitor [--start] [--stop] [-c <ch>]\n");
        return 0;
    }

    // 1. Handle Channel Set
    uint8_t ch = 0;
    if (mon_args.channel->count > 0) {
        ch = (uint8_t)mon_args.channel->ival[0];
        // Note: You could add a setter to wifi_controller here if you want it staged
    }

    // 2. Handle Actions
    if (mon_args.stop->count > 0) {
        wifi_ctl_switch_to_sta(WIFI_BW_HT20);
        printf("Switched to Station Mode.\n");
    }

    if (mon_args.start->count > 0) {
        // If ch is 0, controller uses default/staged
        if (wifi_ctl_switch_to_monitor(ch, WIFI_BW_HT20) == ESP_OK) {
             printf("Monitor Mode Started%s.\n", (ch > 0) ? " (Channel Override)" : "");
        } else {
             printf("Failed to start Monitor Mode.\n");
        }
    }

    if (mon_args.status->count > 0) {
        wifi_ctl_mode_t mode = wifi_ctl_get_mode();
        printf("MONITOR STATUS:\n");
        printf("  Mode:   %s\n", (mode == WIFI_CTL_MODE_MONITOR) ? "MONITOR" : "STATION");
        printf("  Active: Ch %d\n", (mode == WIFI_CTL_MODE_MONITOR) ? wifi_ctl_get_monitor_channel() : 0);
        printf("  Frames: %lu\n", (unsigned long)wifi_ctl_get_monitor_frame_count());
    }

    end_cmd();
    return 0;
}

// ============================================================================
// COMMAND: scan (List Networks)
// ============================================================================
static struct {
    struct arg_lit *help;
    struct arg_end *end;
} scan_args;

static int cmd_scan(int argc, char **argv) {
    int nerrors = arg_parse(argc, argv, (void **)&scan_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, scan_args.end, argv[0]);
        return 1;
    }

    if (scan_args.help->count > 0) {
        printf("Usage: scan\n");
        return 0;
    }

    printf("Starting WiFi Scan...\n");

    // Ensure we are in a mode that supports scanning (STA or NULL)
    wifi_ctl_switch_to_sta(WIFI_BW_HT20);

    wifi_scan_config_t scan_config = {
        .ssid = 0,
        .bssid = 0,
        .channel = 0,
        .show_hidden = true
    };

    esp_err_t err = esp_wifi_scan_start(&scan_config, true); // true = blocking
    if (err != ESP_OK) {
        printf("Scan failed: %s\n", esp_err_to_name(err));
        return 1;
    }

    uint16_t ap_count = 0;
    esp_wifi_scan_get_ap_num(&ap_count);
    printf("Found %d APs:\n", ap_count);

    if (ap_count > 0) {
        wifi_ap_record_t *ap_list = (wifi_ap_record_t *)malloc(sizeof(wifi_ap_record_t) * ap_count);
        if (ap_list) {
            esp_wifi_scan_get_ap_records(&ap_count, ap_list);
            printf("%-32s | %-4s | %-4s | %-3s\n", "SSID", "RSSI", "CH", "Auth");
            printf("----------------------------------------------------------\n");
            for (int i = 0; i < ap_count; i++) {
                printf("%-32s | %-4d | %-4d | %d\n",
                       (char *)ap_list[i].ssid,
                       ap_list[i].rssi,
                       ap_list[i].primary,
                       ap_list[i].authmode);
            }
            free(ap_list);
        }
    }
    return 0;
}

// ============================================================================
// COMMAND: wifi_config
// ============================================================================
static struct {
    struct arg_str *ssid;
    struct arg_str *pass;
    struct arg_str *ip;
    struct arg_lit *dhcp;
    struct arg_lit *help;
    struct arg_end *end;
} wifi_args;

static int cmd_wifi_config(int argc, char **argv) {
    int nerrors = arg_parse(argc, argv, (void **)&wifi_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, wifi_args.end, argv[0]);
        return 1;
    }

    if (wifi_args.help->count > 0) {
        printf("Usage: wifi_config -s <ssid> -p <pass> [-i <ip>] [-d]\n");
        return 0;
    }

    if (wifi_args.ssid->count == 0) {
        printf("Error: SSID is required (-s)\n");
        return 1;
    }

    const char* ssid = wifi_args.ssid->sval[0];
    const char* pass = (wifi_args.pass->count > 0) ? wifi_args.pass->sval[0] : "";
    const char* ip = (wifi_args.ip->count > 0) ? wifi_args.ip->sval[0] : NULL;
    bool dhcp = (wifi_args.dhcp->count > 0);

    printf("Saving WiFi Config to NVS: SSID='%s' DHCP=%d\n", ssid, dhcp);

    // This immediately writes to NVS via wifi_cfg.c
    wifi_cfg_set_credentials(ssid, pass);

    if (ip) {
        char mask[] = "255.255.255.0";
        char gw[32];
        strlcpy(gw, ip, sizeof(gw));
        char *last_dot = strrchr(gw, '.');
        if (last_dot) strcpy(last_dot, ".1"); // Guess Gateway

        wifi_cfg_set_static_ip(ip, mask, gw);
        wifi_cfg_set_dhcp(false);
    } else {
        wifi_cfg_set_dhcp(dhcp);
    }

    printf("Config saved. Rebooting to apply...\n");
    esp_restart();
    return 0;
}

// ============================================================================
// REGISTRATION
// ============================================================================

static void register_iperf_cmd(void) {
    iperf_args.start  = arg_lit0(NULL, "start", "Start");
    iperf_args.stop   = arg_lit0(NULL, "stop", "Stop");
    iperf_args.status = arg_lit0(NULL, "status", "Status");
    iperf_args.save   = arg_lit0(NULL, "save", "Save");
    iperf_args.reload = arg_lit0(NULL, "reload", "Reload");
    iperf_args.ip     = arg_str0("c", "client", "<ip>", "IP");
    iperf_args.port   = arg_int0("p", "port", "<port>", "Port");
    iperf_args.pps    = arg_int0(NULL, "pps", "<n>", "PPS");
    iperf_args.len    = arg_int0(NULL, "len", "<bytes>", "Len");
    iperf_args.burst  = arg_int0(NULL, "burst", "<count>", "Burst");
    iperf_args.help   = arg_lit0("h", "help", "Help");
    iperf_args.end    = arg_end(20);

    const esp_console_cmd_t cmd = { .command = "iperf", .help = "Traffic Gen", .func = &cmd_iperf, .argtable = &iperf_args };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

static void register_monitor_cmd(void) {
    mon_args.start   = arg_lit0(NULL, "start", "Start Monitor");
    mon_args.stop    = arg_lit0(NULL, "stop", "Stop (STA)");
    mon_args.status  = arg_lit0(NULL, "status", "Status");
    mon_args.channel = arg_int0("c", "channel", "<n>", "Channel");
    mon_args.help    = arg_lit0("h", "help", "Help");
    mon_args.end     = arg_end(20);

    const esp_console_cmd_t cmd = { .command = "monitor", .help = "Switch Mode (STA/Monitor)", .func = &cmd_monitor, .argtable = &mon_args };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

static void register_scan_cmd(void) {
    scan_args.help = arg_lit0("h", "help", "Help");
    scan_args.end  = arg_end(1);

    const esp_console_cmd_t cmd = { .command = "scan", .help = "Scan WiFi", .func = &cmd_scan, .argtable = &scan_args };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

static void register_wifi_cmd(void) {
    wifi_args.ssid = arg_str0("s", "ssid", "<ssid>", "SSID");
    wifi_args.pass = arg_str0("p", "password", "<pass>", "Pass");
    wifi_args.ip   = arg_str0("i", "ip", "<ip>", "Static IP");
    wifi_args.dhcp = arg_lit0("d", "dhcp", "DHCP");
    wifi_args.help = arg_lit0("h", "help", "Help");
    wifi_args.end  = arg_end(20);

    const esp_console_cmd_t cmd = { .command = "wifi_config", .help = "Setup WiFi Creds", .func = &cmd_wifi_config, .argtable = &wifi_args };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

void app_console_register_commands(void) {
    register_iperf_cmd();
    register_monitor_cmd();
    register_scan_cmd();
    register_wifi_cmd();
}