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

// Dependencies
#include "iperf.h"
#include "status_led.h"
#include "wifi_monitor.h"
#include "gps_sync.h"

#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
#include "csi_manager.h"
#endif

static const char *TAG = "WIFI_CTL";

static wifi_ctl_mode_t s_current_mode = WIFI_CTL_MODE_STA;
static uint8_t s_monitor_channel_active = 6;
static uint8_t s_monitor_channel_staging = 6;

static bool s_monitor_enabled = false;
static uint32_t s_monitor_frame_count = 0;
static TaskHandle_t s_monitor_stats_task_handle = NULL;

// --- Event Handler ---
static void wifi_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data) {
    if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ESP_LOGI(TAG, "Got IP -> LED Connected");
        status_led_set_state(LED_STATE_CONNECTED);
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        if (s_current_mode == WIFI_CTL_MODE_STA) {
             status_led_set_state(LED_STATE_NO_CONFIG);
        }
    }
}

// ... [Log Collapse / Monitor Callback Logic] ...
static void log_collapse_event(uint32_t nav_duration_us, int rssi, int retry) {
    gps_timestamp_t ts = gps_get_timestamp();
    int64_t now_ms = ts.gps_us / 1000;
    ESP_LOGI(TAG, "COLLAPSE: Time=%" PRId64 "ms, Sync=%d, Dur=%lu us, RSSI=%d, Retry=%d",
             now_ms, ts.synced ? 1 : 0, nav_duration_us, rssi, retry);
}

static void monitor_frame_callback(const wifi_frame_info_t *frame, const uint8_t *payload, uint16_t len) {
    s_monitor_frame_count++;
    if (frame->retry && frame->duration_id > 5000) {
        log_collapse_event((float)frame->duration_id, frame->rssi, frame->retry);
    }
}

static void monitor_stats_task(void *arg) {
    while (1) {
        vTaskDelay(pdMS_TO_TICKS(10000));
        wifi_collapse_stats_t stats;
        if (wifi_monitor_get_stats(&stats) == ESP_OK) {
            ESP_LOGI("MONITOR", "--- Stats: %lu frames, Retry: %.2f%%, Avg NAV: %u us ---",
                     (unsigned long)stats.total_frames, stats.retry_rate, stats.avg_nav);
            if (wifi_monitor_is_collapsed()) ESP_LOGW("MONITOR", "⚠️  COLLAPSE DETECTED! ⚠️");
        }
    }
}

// --- Helper to apply IP settings ---
static void apply_ip_settings(void) {
    esp_netif_t *netif = esp_netif_get_handle_from_ifkey("WIFI_STA_DEF");
    if (!netif) return;

    if (wifi_cfg_get_dhcp()) {
        esp_netif_dhcpc_start(netif);
    } else {
        esp_netif_dhcpc_stop(netif);

        char ip[16], mask[16], gw[16];
        if (wifi_cfg_get_ipv4(ip, mask, gw)) {
            esp_netif_ip_info_t info = {0};
            // API Fix: esp_ip4addr_aton returns uint32_t
            info.ip.addr = esp_ip4addr_aton(ip);
            info.netmask.addr = esp_ip4addr_aton(mask);
            info.gw.addr = esp_ip4addr_aton(gw);

            esp_netif_set_ip_info(netif, &info);
            ESP_LOGI(TAG, "Static IP applied: %s", ip);
        }
    }
}

// ============================================================================
// PUBLIC API IMPLEMENTATION
// ============================================================================

void wifi_ctl_init(void) {
    s_current_mode = WIFI_CTL_MODE_STA;
    s_monitor_enabled = false;
    s_monitor_frame_count = 0;

    // 1. Initialize Network Interface
    esp_netif_create_default_wifi_sta();

    // 2. Apply IP Settings (Static vs DHCP)
    apply_ip_settings();

    // 3. Initialize Wi-Fi Driver
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // 4. Register Events
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED, &wifi_event_handler, NULL, NULL));

    // 5. Configure Storage & Mode
    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    // 6. Apply Saved Config
    if (!wifi_cfg_apply_from_nvs()) {
        ESP_LOGW(TAG, "No saved WiFi config found, driver initialized in defaults.");
        status_led_set_state(LED_STATE_NO_CONFIG);
    } else {
        ESP_LOGI(TAG, "WiFi driver initialized from NVS.");
        status_led_set_state(LED_STATE_WAITING);
        esp_wifi_connect();
    }

    // Load Staging Params
    char mode_ignored[16];
    wifi_cfg_get_mode(mode_ignored, &s_monitor_channel_staging);
    if (s_monitor_channel_staging == 0) s_monitor_channel_staging = 6;
}

// --- Mode Control (Core) ---

esp_err_t wifi_ctl_switch_to_monitor(uint8_t channel, wifi_bandwidth_t bw) {
    if (channel == 0) channel = s_monitor_channel_staging;

    if (s_current_mode == WIFI_CTL_MODE_MONITOR && s_monitor_channel_active == channel) {
        ESP_LOGW(TAG, "Already in monitor mode (Ch %d)", channel);
        return ESP_OK;
    }

    ESP_LOGI(TAG, "Switching to MONITOR MODE (Ch %d)", channel);

    iperf_stop();
    vTaskDelay(pdMS_TO_TICKS(500));
#ifdef CONFIG_ESP_WIFI_CSI_ENABLED
    csi_mgr_disable();
#endif

    esp_wifi_disconnect();
    esp_wifi_stop();
    vTaskDelay(pdMS_TO_TICKS(500));

    esp_wifi_set_mode(WIFI_MODE_NULL);
    if (wifi_monitor_init(channel, monitor_frame_callback) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to init monitor mode");
        return ESP_FAIL;
    }

    esp_wifi_set_bandwidth(WIFI_IF_STA, bw);

    if (wifi_monitor_start() != ESP_OK) {
        ESP_LOGE(TAG, "Failed to start monitor mode");
        return ESP_FAIL;
    }

    s_monitor_enabled = true;
    s_current_mode = WIFI_CTL_MODE_MONITOR;
    s_monitor_channel_active = channel;
    status_led_set_state(LED_STATE_MONITORING);

    if (s_monitor_stats_task_handle == NULL) {
        xTaskCreate(monitor_stats_task, "monitor_stats", 4096, NULL, 5, &s_monitor_stats_task_handle);
    }

    return ESP_OK;
}

esp_err_t wifi_ctl_switch_to_sta(void) {
    if (s_current_mode == WIFI_CTL_MODE_STA) {
        ESP_LOGI(TAG, "Already in STA mode");
        return ESP_OK;
    }

    ESP_LOGI(TAG, "Switching to STA MODE");

    if (s_monitor_stats_task_handle != NULL) {
        vTaskDelete(s_monitor_stats_task_handle);
        s_monitor_stats_task_handle = NULL;
    }

    if (s_monitor_enabled) {
        wifi_monitor_stop();
        s_monitor_enabled = false;
        vTaskDelay(pdMS_TO_TICKS(500));
    }

    esp_wifi_set_mode(WIFI_MODE_STA);
    vTaskDelay(pdMS_TO_TICKS(500));

    esp_wifi_start();
    esp_wifi_connect();

    s_current_mode = WIFI_CTL_MODE_STA;
    status_led_set_state(LED_STATE_WAITING);

    return ESP_OK;
}

// --- Wrappers for cmd_monitor.c ---

void wifi_ctl_monitor_start(int channel) {
    wifi_ctl_switch_to_monitor((uint8_t)channel, WIFI_BW_HT20);
}

void wifi_ctl_stop(void) {
    wifi_ctl_switch_to_sta();
}

void wifi_ctl_start_station(void) {
    wifi_ctl_switch_to_sta();
}

void wifi_ctl_start_ap(void) {
    ESP_LOGW(TAG, "AP Mode not fully implemented, using STA");
    wifi_ctl_switch_to_sta();
}

// --- Settings ---

void wifi_ctl_set_channel(int channel) {
    if (channel < 1 || channel > 14) {
        ESP_LOGE(TAG, "Invalid channel %d", channel);
        return;
    }
    s_monitor_channel_staging = (uint8_t)channel;

    if (s_current_mode == WIFI_CTL_MODE_MONITOR) {
        ESP_LOGI(TAG, "Switching live channel to %d", channel);
        esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
        s_monitor_channel_active = (uint8_t)channel;
    }
}

void wifi_ctl_status(void) {
    const char *mode_str = (s_current_mode == WIFI_CTL_MODE_MONITOR) ? "MONITOR" :
                           (s_current_mode == WIFI_CTL_MODE_AP) ? "AP" : "STATION";

    printf("WiFi Status:\n");
    printf("  Mode:       %s\n", mode_str);
    if (s_current_mode == WIFI_CTL_MODE_MONITOR) {
        printf("  Channel:    %d\n", s_monitor_channel_active);
        printf("  Frames:     %lu\n", (unsigned long)s_monitor_frame_count);
    }
    printf("  Staging Ch: %d\n", s_monitor_channel_staging);
}

// --- Params (NVS) ---

bool wifi_ctl_param_is_unsaved(void) {
    return wifi_cfg_monitor_channel_is_unsaved(s_monitor_channel_staging);
}

void wifi_ctl_param_save(const char *dummy) {
    (void)dummy;
    if (wifi_cfg_set_monitor_channel(s_monitor_channel_staging)) {
        ESP_LOGI(TAG, "Monitor channel (%d) saved to NVS", s_monitor_channel_staging);
    } else {
        ESP_LOGI(TAG, "No changes to save.");
    }
}

void wifi_ctl_param_init(void) {
    char mode_ignored[16];
    uint8_t ch = 0;
    wifi_cfg_get_mode(mode_ignored, &ch);
    if (ch > 0) s_monitor_channel_staging = ch;
    ESP_LOGI(TAG, "Reloaded monitor channel: %d", s_monitor_channel_staging);
}

void wifi_ctl_param_clear(void) {
    wifi_cfg_clear_monitor_channel();
    s_monitor_channel_staging = 6;
    ESP_LOGI(TAG, "Monitor config cleared (Defaulting to Ch 6).");
}

// --- Getters ---

wifi_ctl_mode_t wifi_ctl_get_mode(void) { return s_current_mode; }
int wifi_ctl_get_channel(void) { return s_monitor_channel_active; }

// --- Deprecated ---
static void auto_monitor_task_func(void *arg) {
    uint8_t channel = (uint8_t)(uintptr_t)arg;
    ESP_LOGI(TAG, "Waiting for WiFi connection before switching to monitor mode...");
    while (status_led_get_state() != LED_STATE_CONNECTED) {
        vTaskDelay(pdMS_TO_TICKS(500));
    }
    ESP_LOGI(TAG, "WiFi connected, waiting for GPS sync (2s)...");
    vTaskDelay(pdMS_TO_TICKS(2000));
    wifi_ctl_switch_to_monitor(channel, WIFI_BW_HT20);
    vTaskDelete(NULL);
}

void wifi_ctl_auto_monitor_start(uint8_t channel) {
    xTaskCreate(auto_monitor_task_func, "auto_monitor", 4096, (void*)(uintptr_t)channel, 5, NULL);
}