/*
 * 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 "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_wifi.h"
#include "wifi_cfg.h"
#include "wifi_controller.h"
#include "app_console.h"

// --- Arguments ---
static struct {
    struct arg_str *ssid;
    struct arg_str *password;
    struct arg_end *end;
} connect_args;

static struct {
    struct arg_str *mode; // "sta", "monitor", "ap"
    struct arg_int *channel;
    struct arg_end *end;
} mode_args;

static struct {
    struct arg_int *tx_power;
    struct arg_end *end;
} power_args;

static void print_wifi_usage(void) {
    printf("Usage: wifi <subcommand> [args]\n");
    printf("Subcommands:\n");
    printf("  scan             Scan for networks\n");
    printf("  connect <ssid> [<pass>]   Connect to AP\n");
    printf("  status           Show connection info\n");
    printf("  mode <sta|monitor> [-c <ch>]  Switch Mode\n");
    printf("  power <dBm>      Set TX Power (8-84, 0.25dB units)\n");
}

// --- Handlers ---

static int wifi_do_scan(int argc, char **argv) {
    printf("Scanning...\n");
    wifi_scan_config_t scan_config = {0};
    esp_wifi_scan_start(&scan_config, true); // Block until done

    uint16_t ap_num = 0;
    esp_wifi_scan_get_ap_num(&ap_num);

    wifi_ap_record_t *ap_list = (wifi_ap_record_t *)malloc(ap_num * sizeof(wifi_ap_record_t));
    ESP_ERROR_CHECK(esp_wifi_scan_get_ap_records(&ap_num, ap_list));

    printf("Found %d APs:\n", ap_num);
    printf("%-32s | %-4s | %-4s | %-18s\n", "SSID", "RSSI", "CH", "BSSID");
    for (int i = 0; i < ap_num; i++) {
        printf("%-32s | %-4d | %-4d | %02x:%02x:%02x:%02x:%02x:%02x\n",
               ap_list[i].ssid, ap_list[i].rssi, ap_list[i].primary,
               ap_list[i].bssid[0], ap_list[i].bssid[1], ap_list[i].bssid[2],
               ap_list[i].bssid[3], ap_list[i].bssid[4], ap_list[i].bssid[5]);
    }
    free(ap_list);
    return 0;
}

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

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

    // Ensure we are in STA mode first
    wifi_ctl_mode_t current_mode = wifi_ctl_get_mode();
    if (current_mode != WIFI_CTL_MODE_STA) {
        printf("Switching to Station Mode first...\n");
        wifi_ctl_switch_to_sta(); // Fixed: Removed argument
    }

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

    // Save to NVS
    wifi_cfg_set_ssid(ssid);
    wifi_cfg_set_password(pass);

    // Apply
    wifi_config_t wifi_config = {0};
    strncpy((char *)wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid) - 1);
    wifi_config.sta.ssid[sizeof(wifi_config.sta.ssid) - 1] = '\0';
    strncpy((char *)wifi_config.sta.password, pass, sizeof(wifi_config.sta.password) - 1);
    wifi_config.sta.password[sizeof(wifi_config.sta.password) - 1] = '\0';

    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_connect());

    return 0;
}

static int wifi_do_status(int argc, char **argv) {
    wifi_ctl_status();
    return 0;
}

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

    const char *mode_str = mode_args.mode->sval[0];
    int channel = (mode_args.channel->count > 0) ? mode_args.channel->ival[0] : 0;

    if (strcmp(mode_str, "sta") == 0) {
        wifi_ctl_switch_to_sta(); // Fixed: Removed argument
    } else if (strcmp(mode_str, "monitor") == 0) {
        wifi_ctl_switch_to_monitor(channel, WIFI_BW_HT20);
    } else {
        printf("Unknown mode '%s'. Use 'sta' or 'monitor'.\n", mode_str);
        return 1;
    }
    return 0;
}

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

    int pwr = power_args.tx_power->ival[0];
    esp_err_t err = esp_wifi_set_max_tx_power(pwr);
    if (err == ESP_OK) {
        printf("TX Power set to %d (approx %.2f dBm)\n", pwr, pwr * 0.25);
    } else {
        printf("Failed to set TX power: %s\n", esp_err_to_name(err));
    }
    return 0;
}

static int cmd_wifi(int argc, char **argv) {
    if (argc < 2) {
        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], "status") == 0)  return wifi_do_status(argc - 1, &argv[1]);
    if (strcmp(argv[1], "mode") == 0)    return wifi_do_mode(argc - 1, &argv[1]);
    if (strcmp(argv[1], "power") == 0)   return wifi_do_power(argc - 1, &argv[1]);

    if (strcmp(argv[1], "help") == 0 || strcmp(argv[1], "--help") == 0) {
        print_wifi_usage();
        return 0;
    }

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

void register_wifi_cmd(void) {
    // Connect Args
    connect_args.ssid = arg_str1(NULL, NULL, "<ssid>", "SSID");
    connect_args.password = arg_str0(NULL, NULL, "<pass>", "Password");
    connect_args.end = arg_end(2);

    // Mode Args
    mode_args.mode = arg_str1(NULL, NULL, "<mode>", "sta | monitor");
    mode_args.channel = arg_int0("c", "channel", "<n>", "Channel (Monitor only)");
    mode_args.end = arg_end(2);

    // Power Args
    power_args.tx_power = arg_int1(NULL, NULL, "<dBm>", "Power (8-84)");
    power_args.end = arg_end(1);

    const esp_console_cmd_t cmd = {
        .command = "wifi",
        .help = "Wi-Fi Utils: scan, connect, mode, status, power",
        .hint = "<subcommand>",
        .func = &cmd_wifi,
    };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}