/*
 * 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 for Sub-Command Handlers ---
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);

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

static int cmd_wifi(int argc, char **argv) {
    if (argc < 2) {
        printf("Usage: wifi <subcommand> [args]\n");
        printf("Subcommands:\n");
        printf("  link        Show connection status\n");
        printf("  scan        Scan for networks\n");
        printf("  connect     Connect to an AP\n");
        printf("  disconnect  Disconnect from AP\n");
        return 0;
    }

    // Dispatch based on subcommand
    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], "disconnect") == 0) return wifi_do_disconnect(argc - 1, &argv[1]);
    if (strcmp(argv[1], "link") == 0)       return wifi_do_link(argc - 1, &argv[1]);

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

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

static int wifi_do_link(int argc, char **argv) {
    // Parse args mainly to support --help
    if (arg_parse(argc, argv, (void **)&link_args) == 0) {
        if (link_args.help->count > 0) {
            printf("Usage: wifi link\nDescription: Show connection status (RSSI, BSSID, SSID, IP).\n");
            return 0;
        }
    }

    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\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));
            printf("\tMask:   " IPSTR "\n", IP2STR(&ip.netmask));
            printf("\tGW:     " IPSTR "\n", IP2STR(&ip.gw));
        }
    } else {
        printf("Not connected.\n");
    }
    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) {
    if (arg_parse(argc, argv, (void **)&scan_args) == 0) {
        if (scan_args.help->count > 0) {
            printf("Usage: wifi scan\n");
            return 0;
        }
    }

    // Prevent scanning in Monitor Mode (Hardware limitation)
    if (wifi_ctl_get_mode() == WIFI_CTL_MODE_MONITOR) {
        printf("Error: Cannot scan while in Monitor Mode.\n");
        printf("Run 'monitor --stop' first, or 'wifi disconnect'.\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);

    // Smart Retry: If busy connecting, force disconnect momentarily
    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 for scan list.\n");
        return 1;
    }

    ESP_ERROR_CHECK(esp_wifi_scan_get_ap_records(&ap_count, ap_info));

    printf("Found %u unique BSSIDs\n", ap_count);
    printf("BSS                             | SSID                             | RSSI | Ch | Auth\n");
    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;
            default: break;
        }

        printf("%02x:%02x:%02x:%02x:%02x:%02x | %-32.32s | %4d | %2d | %s\n",
               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].ssid, ap_info[i].rssi, ap_info[i].primary, authmode);
    }
    printf("-------------------------------------------------------------------------------------\n");
    free(ap_info);
    return 0;
}

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

static int wifi_do_connect(int argc, char **argv) {
    // Re-init pointers to be safe with static structs
    connect_args.ssid = arg_str1(NULL, NULL, "<ssid>", "SSID");
    connect_args.password = arg_str0(NULL, NULL, "<pass>", "Password");
    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]\n");
        return 0;
    }

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

    printf("Connecting to '%s'...\n", ssid);

    // Save to NVS
    wifi_cfg_set_credentials(ssid, pass);
    wifi_cfg_set_dhcp(true); // Default to DHCP for simple connect

    // Trigger Reconnect
    // The most robust way to switch networks completely is to restart,
    // ensuring the WiFi Controller state machine resets cleanly from NVS.
    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();

    // Optionally clear credentials if you want "disconnect" to mean "forget"
    // wifi_cfg_clear_credentials();
    // printf("Credentials forgotten.\n");

    return 0;
}

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

void register_wifi_cmd(void) {
    const esp_console_cmd_t cmd = {
        .command = "wifi",
        .help = "Wi-Fi Configuration Tool (iw style)",
        .func = &cmd_wifi,
        .argtable = NULL // Handled internally by sub-dispatchers
    };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}