#include "wifi_manager.h"
#include "pico/cyw43_arch.h"
#include "pico/multicore.h"
#include "pico/flash.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);
    }
}

// Helper struct for flash operations
struct FlashWriteParams {
    const uint8_t* data;
    size_t size;
};

// Actual flash operation to be executed safely
static void do_flash_write(void *param) {
    FlashWriteParams *p = (FlashWriteParams*)param;
    flash_range_erase(FLASH_TARGET_OFFSET, FLASH_SECTOR_SIZE);
    flash_range_program(FLASH_TARGET_OFFSET, p->data, p->size);
}

int 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';

    FlashWriteParams params = { (const uint8_t*)&creds, sizeof(creds) };
    int rc = flash_safe_execute(do_flash_write, &params, 1000);
    
    if (rc == PICO_OK) {
        printf("WiFi credentials saved to flash.\n");
    } else {
        printf("Failed to save WiFi credentials: %d\n", rc);
    }
    return rc;
}

bool wifi_try_auto_connect() {
    // On RP2040/RP2350, flash is memory-mapped at XIP_BASE (0x10000000).
    // We can read from it directly like a normal pointer without special API calls.
    // The hardware XIP controller handles fetching data from the flash chip.
    const WifiCreds* creds = (const WifiCreds*)(XIP_BASE + FLASH_TARGET_OFFSET);
    
    // Check for the magic number to verify valid data exists.
    // Erased flash reads as 0xFFFFFFFF, so this check fails if no data was saved.
    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]));
}
