#include "wifi_scan.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "cJSON.h"
#include <string.h>
#include <stdlib.h>

#define TAG "wifi_scan"
#define MAX_NETWORKS 100
#define MAX_JSON_SIZE 16384

// Static buffer for JSON results
static char scan_results_json[MAX_JSON_SIZE];
static int network_count = 0;

// Helper function to get security type string
static const char* get_security_type(wifi_auth_mode_t authmode) {
    switch (authmode) {
        case WIFI_AUTH_OPEN:
            return "Open";
        case WIFI_AUTH_WEP:
            return "WEP";
        case WIFI_AUTH_WPA_PSK:
            return "WPA PSK";
        case WIFI_AUTH_WPA2_PSK:
            return "WPA2 PSK";
        case WIFI_AUTH_WPA_WPA2_PSK:
            return "WPA/WPA2 PSK";
        case WIFI_AUTH_WPA3_PSK:
            return "WPA3 PSK";
        case WIFI_AUTH_WPA2_WPA3_PSK:
            return "WPA2/WPA3 PSK";
        default:
            return "Unknown";
    }
}

// Helper function to determine band from channel
static const char* get_band(uint8_t channel) {
    if (channel >= 1 && channel <= 13) {
        return "2.4GHz";
    }
    if ((channel >= 36 && channel <= 144) || (channel >= 149 && channel <= 165)) {
        return "5GHz";
    }
    return "Unknown";
}

// Helper function to get PHY mode string
static const char* get_phy_mode(const wifi_ap_record_t *ap) {
    if (ap->phy_11n) {
        return "802.11n";
    }
    if (ap->primary > 14) {
        return "802.11a";
    }
    return "802.11b/g";
}

int wifi_scan_networks(void) {
    ESP_LOGI(TAG, "Starting dual-band WiFi scan...");
    
    // Clear previous results
    network_count = 0;
    scan_results_json[0] = '\0';
    
    // Configure scan
    wifi_scan_config_t scan_config = {
        .ssid = NULL,
        .bssid = NULL,
        .channel = 0,  // Scan all channels
        .show_hidden = true,
        .scan_type = WIFI_SCAN_TYPE_ACTIVE,
        .scan_time.active.min = 50,
        .scan_time.active.max = 100,
        .scan_time.passive = 100
    };
    
    // Start scan
    esp_err_t err = esp_wifi_scan_start(&scan_config, true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "WiFi scan failed: %s", esp_err_to_name(err));
        return -1;
    }
    
    // Get number of APs found
    uint16_t ap_count = 0;
    err = esp_wifi_scan_get_ap_num(&ap_count);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to get AP count: %s", esp_err_to_name(err));
        return -1;
    }
    
    if (ap_count == 0) {
        ESP_LOGI(TAG, "No networks found");
        strcpy(scan_results_json, "[]");
        network_count = 0;
        return 0;
    }
    
    // Limit to MAX_NETWORKS
    if (ap_count > MAX_NETWORKS) {
        ap_count = MAX_NETWORKS;
    }
    
    // Allocate memory for AP records
    wifi_ap_record_t *ap_records = malloc(sizeof(wifi_ap_record_t) * ap_count);
    if (!ap_records) {
        ESP_LOGE(TAG, "Failed to allocate memory for scan results");
        return -1;
    }
    
    // Get AP records
    err = esp_wifi_scan_get_ap_records(&ap_count, ap_records);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to get AP records: %s", esp_err_to_name(err));
        free(ap_records);
        return -1;
    }
    
    // Create JSON array
    cJSON *networks_array = cJSON_CreateArray();
    if (!networks_array) {
        ESP_LOGE(TAG, "Failed to create JSON array");
        free(ap_records);
        return -1;
    }
    
    // Process each network
    for (int i = 0; i < ap_count; i++) {
        cJSON *network = cJSON_CreateObject();
        if (!network) {
            continue;
        }
        
        // SSID (handle hidden networks)
        const char *ssid_str = (ap_records[i].ssid[0] == 0) ? "(hidden)" : (char*)ap_records[i].ssid;
        cJSON_AddStringToObject(network, "ssid", ssid_str);
        
        // BSSID
        char bssid_str[18];
        snprintf(bssid_str, sizeof(bssid_str), "%02x:%02x:%02x:%02x:%02x:%02x",
                 ap_records[i].bssid[0], ap_records[i].bssid[1], ap_records[i].bssid[2],
                 ap_records[i].bssid[3], ap_records[i].bssid[4], ap_records[i].bssid[5]);
        cJSON_AddStringToObject(network, "bssid", bssid_str);
        
        // Channel
        cJSON_AddNumberToObject(network, "channel", ap_records[i].primary);
        
        // Band
        cJSON_AddStringToObject(network, "band", get_band(ap_records[i].primary));
        
        // RSSI
        cJSON_AddNumberToObject(network, "rssi", ap_records[i].rssi);
        
        // Security
        cJSON_AddStringToObject(network, "security", get_security_type(ap_records[i].authmode));
        
        // PHY mode
        cJSON_AddStringToObject(network, "phy_mode", get_phy_mode(&ap_records[i]));
        
        // Second channel (if using 40MHz)
        if (ap_records[i].second) {
            cJSON_AddNumberToObject(network, "second_channel", ap_records[i].second);
        }
        
        // Hidden network flag
        cJSON_AddBoolToObject(network, "is_hidden", (ap_records[i].ssid[0] == 0));
        
        cJSON_AddItemToArray(networks_array, network);
    }
    
    // Generate JSON string
    char *json_string = cJSON_PrintUnformatted(networks_array);
    if (json_string) {
        size_t json_len = strlen(json_string);
        
        // Check if JSON is too large and warn
        if (json_len >= MAX_JSON_SIZE) {
            ESP_LOGW(TAG, "JSON too large (%zu >= %d), truncating", json_len, MAX_JSON_SIZE);
        }
        
        // Safe copy to static buffer with guaranteed null termination
        size_t copy_len = (json_len < MAX_JSON_SIZE - 1) ? json_len : MAX_JSON_SIZE - 1;
        memcpy(scan_results_json, json_string, copy_len);
        scan_results_json[copy_len] = '\0';
        
        free(json_string);
        network_count = ap_count;
    } else {
        ESP_LOGE(TAG, "Failed to generate JSON string");
        cJSON_Delete(networks_array);
        free(ap_records);
        return -1;
    }
    
    cJSON_Delete(networks_array);
    free(ap_records);
    
    ESP_LOGI(TAG, "Scan completed, found %d networks", network_count);
    return 0;
}

const char* wifi_scan_get_results_json(void) {
    if (network_count == 0 && scan_results_json[0] == '\0') {
        return "[]";
    }
    return scan_results_json;
}

int wifi_scan_get_count(void) {
    return network_count;
}