#include "wifi_manager.h"
#include "pico/cyw43_arch.h"
#include "hardware/flash.h"
#include "hardware/sync.h"
#include <stdio.h>
#include <string.h>

// Use the last sector of the 2MB flash
#define FLASH_TARGET_OFFSET (2 * 1024 * 1024 - FLASH_SECTOR_SIZE)
#define WIFI_CREDS_MAGIC 0x42574946 // 'BWIF'

struct WifiCreds {
    uint32_t magic;
    char ssid[33];
    char password[64];
};

static volatile bool g_wifi_initialized = false;
static bool g_wifi_connected = false;

bool wifi_init() {
    if (g_wifi_initialized) return true;

    if (cyw43_arch_init()) {
        printf("WiFi init failed!\n");
        return false;
    }

    cyw43_arch_enable_sta_mode();
    g_wifi_initialized = true;
    printf("WiFi initialized on Core %u\n", get_core_num());
    return true;
}

static int scan_result(void *env, const cyw43_ev_scan_result_t *result) {
    if (result) {
        printf("SSID: %-32s RSSI: %4d Auth: %u\n", result->ssid, result->rssi, result->auth_mode);
    }
    return 0;
}

void wifi_scan() {
    // Wait for initialization
    int timeout = 500; // 5 seconds
    while (!g_wifi_initialized && timeout > 0) {
        sleep_ms(10);
        timeout--;
    }

    if (!g_wifi_initialized) {
        printf("WiFi not initialized!\n");
        return;
    }
    
    printf("Starting WiFi scan...\n");
    cyw43_wifi_scan_options_t scan_options = {0};
    int err = cyw43_wifi_scan(&cyw43_state, &scan_options, NULL, scan_result);
    if (err != 0) {
        printf("Failed to start scan: %d\n", err);
    }
}

void wifi_save_credentials(const char* ssid, const char* password) {
    WifiCreds creds;
    creds.magic = WIFI_CREDS_MAGIC;
    strncpy(creds.ssid, ssid, sizeof(creds.ssid) - 1);
    creds.ssid[sizeof(creds.ssid) - 1] = '\0';
    strncpy(creds.password, password, sizeof(creds.password) - 1);
    creds.password[sizeof(creds.password) - 1] = '\0';

    uint32_t ints = save_and_disable_interrupts();
    flash_range_erase(FLASH_TARGET_OFFSET, FLASH_SECTOR_SIZE);
    flash_range_program(FLASH_TARGET_OFFSET, (const uint8_t*)&creds, sizeof(creds));
    restore_interrupts(ints);
    
    printf("WiFi credentials saved to flash.\n");
}

bool wifi_try_auto_connect() {
    const WifiCreds* creds = (const WifiCreds*)(XIP_BASE + FLASH_TARGET_OFFSET);
    
    if (creds->magic != WIFI_CREDS_MAGIC) {
        printf("No saved WiFi credentials found.\n");
        return false;
    }
    
    printf("Found saved credentials for SSID: %s\n", creds->ssid);
    return wifi_connect(creds->ssid, creds->password);
}

bool wifi_connect(const char* ssid, const char* password) {
    // Wait for initialization
    int timeout = 500; // 5 seconds
    while (!g_wifi_initialized && timeout > 0) {
        sleep_ms(10);
        timeout--;
    }

    if (!g_wifi_initialized) {
        printf("WiFi not initialized!\n");
        return false;
    }

    printf("Connecting to WiFi: %s...\n", ssid);
    
    // Connect in blocking mode for simplicity
    if (cyw43_arch_wifi_connect_timeout_ms(ssid, password, CYW43_AUTH_WPA2_AES_PSK, 30000)) {
        printf("WiFi connection failed!\n");
        g_wifi_connected = false;
        return false;
    }

    printf("WiFi connected!\n");
    g_wifi_connected = true;
    return true;
}

bool wifi_is_connected() {
    return g_wifi_connected;
}

const char* wifi_get_ip() {
    if (!g_wifi_connected) return "Disconnected";
    return ip4addr_ntoa(netif_ip4_addr(&cyw43_state.netif[0]));
}