/*
 * cmd_wifi.c
 *
 * Copyright (c) 2025 Umber Networks & Robert McMahon
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */
#include <stdio.h>
#include <string.h>
#include "esp_log.h"
#include "esp_console.h"
#include "argtable3/argtable3.h"
#include "esp_wifi.h"
#include "esp_netif.h"
#include "wifi_cfg.h"
#include "wifi_controller.h"
#include "app_console.h"

// --- Forward Declarations ---
static int wifi_do_scan(int argc, char **argv);
static int wifi_do_connect(int argc, char **argv);
static int wifi_do_disconnect(int argc, char **argv);
static int wifi_do_link(int argc, char **argv);
static int wifi_do_power(int argc, char **argv);
static int wifi_do_mode(int argc, char **argv);

// ============================================================================
// COMMAND: wifi (Dispatcher)
// ============================================================================

static void print_wifi_usage(void) {
    printf("Usage: wifi <subcommand> [args]\n");
    printf("Subcommands:\n");
    printf("  scan        Scan for networks\n");
    printf("  connect     Connect to an AP (alias: join)\n");
    printf("  disconnect  Disconnect from AP\n");
    printf("  status      Show connection status (alias: link)\n");
    printf("  power       Set power save (on/off)\n");
    printf("  mode        Set mode (sta/monitor)\n");
    printf("\nType 'wifi <subcommand> --help' for details.\n");
}

static int cmd_wifi(int argc, char **argv) {
    if (argc < 2 || strcmp(argv[1], "help") == 0 || strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-h") == 0) {
        print_wifi_usage();
        return 0;
    }

    if (strcmp(argv[1], "scan") == 0)       return wifi_do_scan(argc - 1, &argv[1]);
    if (strcmp(argv[1], "connect") == 0)    return wifi_do_connect(argc - 1, &argv[1]);
    if (strcmp(argv[1], "join") == 0)       return wifi_do_connect(argc - 1, &argv[1]);
    if (strcmp(argv[1], "disconnect") == 0) return wifi_do_disconnect(argc - 1, &argv[1]);
    if (strcmp(argv[1], "status") == 0)     return wifi_do_link(argc - 1, &argv[1]);
    if (strcmp(argv[1], "link") == 0)       return wifi_do_link(argc - 1, &argv[1]);
    if (strcmp(argv[1], "power") == 0)      return wifi_do_power(argc - 1, &argv[1]);
    if (strcmp(argv[1], "mode") == 0)       return wifi_do_mode(argc - 1, &argv[1]);

    printf("Unknown subcommand '%s'.\n", argv[1]);
    print_wifi_usage();
    return 1;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi status / link
// ----------------------------------------------------------------------------
static struct {
    struct arg_lit *help;
    struct arg_end *end;
} link_args;

static int wifi_do_link(int argc, char **argv) {
    link_args.help = arg_lit0("h", "help", "Help");
    link_args.end = arg_end(1);

    if (arg_parse(argc, argv, (void **)&link_args) == 0) {
        if (link_args.help->count > 0) {
            printf("Usage: wifi status\nDescription: Show connection status.\n");
            goto exit;
        }
    }

    wifi_ap_record_t ap_info;
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        printf("Connected to %02x:%02x:%02x:%02x:%02x:%02x (on esp32)\n",
               ap_info.bssid[0], ap_info.bssid[1], ap_info.bssid[2],
               ap_info.bssid[3], ap_info.bssid[4], ap_info.bssid[5]);
        printf("\tSSID:   %s\n", ap_info.ssid);
        printf("\tfreq:   %d (Primary)\n", ap_info.primary);
        printf("\tsignal: %d dBm\n", ap_info.rssi);

        esp_netif_t *netif = esp_netif_get_handle_from_ifkey("WIFI_STA_DEF");
        esp_netif_ip_info_t ip;
        if (netif && esp_netif_get_ip_info(netif, &ip) == ESP_OK) {
            printf("\tIP:     " IPSTR "\n", IP2STR(&ip.ip));
        }
    } else {
        printf("Not connected.\n");
    }

    wifi_ps_type_t ps_type;
    if (esp_wifi_get_ps(&ps_type) == ESP_OK) {
        printf("\tPower Save: %s\n", (ps_type == WIFI_PS_NONE) ? "off" : "on");
    }

exit:
    return 0;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi power
// ----------------------------------------------------------------------------
static struct {
    struct arg_lit *on;
    struct arg_lit *off;
    struct arg_lit *help;
    struct arg_end *end;
} power_args;

static int wifi_do_power(int argc, char **argv) {
    power_args.on   = arg_lit0(NULL, "on", "Enable Power Save");
    power_args.off  = arg_lit0(NULL, "off", "Disable Power Save");
    power_args.help = arg_lit0("h", "help", "Help");
    power_args.end  = arg_end(1);

    int nerrors = arg_parse(argc, argv, (void **)&power_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, power_args.end, argv[0]);
        return 1;
    }
    if (power_args.help->count > 0) {
        printf("Usage: wifi power [on|off]\n");
        return 0;
    }
    if (power_args.on->count > 0) {
        esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
        printf("Power save enabled (WIFI_PS_MIN_MODEM).\n");
    } else if (power_args.off->count > 0) {
        esp_wifi_set_ps(WIFI_PS_NONE);
        printf("Power save disabled (WIFI_PS_NONE).\n");
    } else {
        wifi_ps_type_t ps_type;
        esp_wifi_get_ps(&ps_type);
        printf("Current Power Save Mode: %s\n", (ps_type == WIFI_PS_NONE) ? "OFF" : "ON");
    }
    return 0;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi mode
// ----------------------------------------------------------------------------
static struct {
    struct arg_lit *sta;
    struct arg_lit *monitor;
    struct arg_lit *help;
    struct arg_end *end;
} mode_args;

static int wifi_do_mode(int argc, char **argv) {
    mode_args.sta     = arg_lit0(NULL, "sta", "Station Mode");
    mode_args.monitor = arg_lit0(NULL, "monitor", "Monitor Mode");
    mode_args.help    = arg_lit0("h", "help", "Help");
    mode_args.end     = arg_end(1);

    int nerrors = arg_parse(argc, argv, (void **)&mode_args);
    if (nerrors > 0) {
        arg_print_errors(stderr, mode_args.end, argv[0]);
        return 1;
    }
    if (mode_args.help->count > 0) {
        printf("Usage: wifi mode [sta|monitor]\n");
        return 0;
    }
    if (mode_args.sta->count > 0) {
        wifi_ctl_switch_to_sta(WIFI_BW_HT20);
        printf("Switched to STATION mode.\n");
    } else if (mode_args.monitor->count > 0) {
        wifi_ctl_switch_to_monitor(0, WIFI_BW_HT20);
        printf("Switched to MONITOR mode.\n");
    } else {
        wifi_ctl_mode_t mode = wifi_ctl_get_mode();
        printf("Current Mode: %s\n", (mode == WIFI_CTL_MODE_MONITOR) ? "MONITOR" : "STATION");
    }
    return 0;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi scan
// ----------------------------------------------------------------------------
static struct {
    struct arg_lit *help;
    struct arg_end *end;
} scan_args;

static int wifi_do_scan(int argc, char **argv) {
    scan_args.help = arg_lit0("h", "help", "Help");
    scan_args.end  = arg_end(1);

    if (arg_parse(argc, argv, (void **)&scan_args) == 0) {
        if (scan_args.help->count > 0) {
            printf("Usage: wifi scan\n");
            return 0;
        }
    }
    if (wifi_ctl_get_mode() == WIFI_CTL_MODE_MONITOR) {
        printf("Error: Cannot scan while in Monitor Mode.\n");
        return 1;
    }
    wifi_scan_config_t scan_config = {0};
    scan_config.show_hidden = true;
    printf("Scanning...\n");
    esp_err_t err = esp_wifi_scan_start(&scan_config, true);
    if (err == ESP_ERR_WIFI_STATE) {
        printf("WARN: Interface busy. Forcing disconnect for scan...\n");
        esp_wifi_disconnect();
        vTaskDelay(pdMS_TO_TICKS(100));
        err = esp_wifi_scan_start(&scan_config, true);
    }
    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);
    if (ap_count == 0) {
        printf("No networks found.\n");
        return 0;
    }
    wifi_ap_record_t *ap_info = (wifi_ap_record_t *)malloc(sizeof(wifi_ap_record_t) * ap_count);
    if (!ap_info) {
        printf("Out of memory.\n");
        return 1;
    }
    ESP_ERROR_CHECK(esp_wifi_scan_get_ap_records(&ap_count, ap_info));
    printf("Found %u unique BSSIDs\n", ap_count);

    // FORMATTING: Using exact width specifiers for Header and Data to ensure alignment
    // SSID: 32 chars | BSSID: 17 chars | RSSI: 4 chars | Ch: 3 chars | Auth
    printf("%-32s | %-17s | %4s | %3s | %s\n", "SSID", "BSSID", "RSSI", "Ch", "Auth");
    printf("--------------------------------------------------------------------------------------\n");

    for (int i = 0; i < ap_count; i++) {
        char *authmode = "UNK";
        switch (ap_info[i].authmode) {
            case WIFI_AUTH_OPEN: authmode = "OPEN"; break;
            case WIFI_AUTH_WEP: authmode = "WEP"; break;
            case WIFI_AUTH_WPA_PSK: authmode = "WPA"; break;
            case WIFI_AUTH_WPA2_PSK: authmode = "WPA2"; break;
            case WIFI_AUTH_WPA_WPA2_PSK: authmode = "WPA/2"; break;
            case WIFI_AUTH_WPA3_PSK: authmode = "WPA3"; break;
            case WIFI_AUTH_WPA2_WPA3_PSK: authmode = "W2/W3"; break;
            default: break;
        }

        printf("%-32.32s | %02x:%02x:%02x:%02x:%02x:%02x | %4d | %3d | %s\n",
               ap_info[i].ssid,
               ap_info[i].bssid[0], ap_info[i].bssid[1], ap_info[i].bssid[2],
               ap_info[i].bssid[3], ap_info[i].bssid[4], ap_info[i].bssid[5],
               ap_info[i].rssi, ap_info[i].primary, authmode);
    }
    printf("--------------------------------------------------------------------------------------\n");
    free(ap_info);
    return 0;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi connect / join
// ----------------------------------------------------------------------------
static struct {
    struct arg_str *ssid;
    struct arg_str *password;
    struct arg_str *bssid;
    struct arg_lit *help;
    struct arg_end *end;
} connect_args;

static int wifi_do_connect(int argc, char **argv) {
    connect_args.ssid     = arg_str1(NULL, NULL, "<ssid>", "SSID");
    connect_args.password = arg_str0(NULL, NULL, "<pass>", "Password");
    connect_args.bssid    = arg_str0("b", "bssid", "<xx:xx...>", "Lock BSSID");
    connect_args.help     = arg_lit0("h", "help", "Help");
    connect_args.end      = arg_end(2);

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

    if (connect_args.help->count > 0) {
        printf("Usage: wifi connect <ssid> [password] [-b <bssid>]\n");
        return 0;
    }

    const char *ssid = connect_args.ssid->sval[0];
    const char *pass = (connect_args.password->count > 0) ? connect_args.password->sval[0] : "";
    const char *bssid = (connect_args.bssid->count > 0) ? connect_args.bssid->sval[0] : "";

    printf("Connecting to '%s' (BSSID: %s)...\n", ssid, bssid[0] ? bssid : "Any");

    wifi_cfg_set_credentials(ssid, pass, bssid);
    wifi_cfg_set_dhcp(true);

    printf("Credentials saved. Rebooting to apply...\n");
    esp_restart();

    return 0;
}

// ----------------------------------------------------------------------------
// Sub-command: wifi disconnect
// ----------------------------------------------------------------------------
static int wifi_do_disconnect(int argc, char **argv) {
    printf("Disconnecting...\n");
    esp_wifi_disconnect();
    return 0;
}

// ----------------------------------------------------------------------------
// Registration
// ----------------------------------------------------------------------------

void register_wifi_cmd(void) {
    const esp_console_cmd_t cmd = {
        .command = "wifi",
        .help = "Wi-Fi Tool: scan, connect, status, power, mode",
        .hint = "<subcommand> [args]",
        .func = &cmd_wifi,
        .argtable = NULL
    };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}