#include "ble_wifi.h"
#include <string.h>
#include "esp_event.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_wifi.h"
#include "esp_wifi_default.h"
#include "freertos/FreeRTOS.h"
#include "freertos/event_groups.h"
#include "freertos/task.h"

#define CONNECTED BIT0
#define DISCONNECTED BIT1
#define STOPPED BIT2
static const char *TAG = "ble_wifi";
static esp_netif_t *station;
static EventGroupHandle_t events;
static esp_event_handler_instance_t wifi_handler, ip_handler;
static bool driver_ready, prepared, started, stopping;

static void wifi_event(void *arg, esp_event_base_t base, int32_t id, void *data)
{
    (void)arg;
    if (base == IP_EVENT && id == IP_EVENT_STA_GOT_IP) {
        xEventGroupSetBits(events, CONNECTED);
        ESP_LOGI(TAG, "Wi-Fi obtained IP");
    } else if (base == WIFI_EVENT && id == WIFI_EVENT_STA_DISCONNECTED) {
        xEventGroupClearBits(events, CONNECTED);
        xEventGroupSetBits(events, DISCONNECTED);
        const wifi_event_sta_disconnected_t *event = data;
        // Never log the SSID, password, binding token, or raw event payload.
        ESP_LOGW(TAG, "Wi-Fi disconnected; reason=%u", event ? event->reason : 0);
    } else if (base == WIFI_EVENT && id == WIFI_EVENT_STA_STOP) {
        xEventGroupClearBits(events, CONNECTED);
        xEventGroupSetBits(events, STOPPED);
    }
}

esp_err_t ble_wifi_prepare(void)
{
    if (prepared) return ESP_OK;
    // Retain each successfully initialized resource across retries. In
    // particular an esp_wifi_init failure must not recreate WIFI_STA or leak
    // another event group on the next screen confirmation.
    if (!events) {
        events = xEventGroupCreate();
        if (!events) return ESP_ERR_NO_MEM;
    }
    if (!station) {
        // Do not adopt/destroy an interface owned by the normal voice flow.
        if (esp_netif_get_handle_from_ifkey("WIFI_STA")) return ESP_ERR_INVALID_STATE;
        esp_netif_config_t config = ESP_NETIF_DEFAULT_WIFI_STA();
        station = esp_netif_new(&config);
        if (!station) return ESP_ERR_NO_MEM;
        // The convenience create_default_wifi_sta() aborts on allocation or
        // duplicate-key failure. Use its checked building blocks instead.
        esp_err_t err = esp_netif_attach_wifi_station(station);
        if (err == ESP_OK) err = esp_wifi_set_default_wifi_sta_handlers();
        if (err != ESP_OK) {
            esp_netif_destroy_default_wifi(station);
            station = NULL;
            return err;
        }
    }
    if (!driver_ready) {
        wifi_init_config_t config = WIFI_INIT_CONFIG_DEFAULT();
        // Provisioning sends a few small HTTP requests, not a voice stream.
        // Static RX/TX buffers require internal RAM even with 8 MB PSRAM.
        // Keep normal voice-mode sdkconfig performance settings unchanged.
        config.static_rx_buf_num = 4;
        config.dynamic_rx_buf_num = 8;
        config.static_tx_buf_num = 4;
        config.cache_tx_buf_num = 4;
        config.ampdu_rx_enable = 0;
        config.ampdu_tx_enable = 0;
        config.amsdu_tx_enable = 0;
        ESP_LOGI(TAG, "Preparing Wi-Fi: internal_free=%u largest=%u rx=%d tx=%d",
                 (unsigned)heap_caps_get_free_size(MALLOC_CAP_INTERNAL),
                 (unsigned)heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL),
                 config.static_rx_buf_num, config.static_tx_buf_num);
        esp_err_t err = esp_wifi_init(&config);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "Wi-Fi driver init failed: %s", esp_err_to_name(err));
            return err;
        }
        driver_ready = true;
    }
    esp_err_t err;
    if (!wifi_handler) {
        err = esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID,
                                                  wifi_event, NULL, &wifi_handler);
        if (err != ESP_OK) return err;
    }
    if (!ip_handler) {
        err = esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                                  wifi_event, NULL, &ip_handler);
        if (err != ESP_OK) return err;
    }
    err = esp_wifi_set_storage(WIFI_STORAGE_RAM);
    if (err == ESP_OK) err = esp_wifi_set_mode(WIFI_MODE_STA);
    if (err != ESP_OK) return err;
    prepared = true;
    ESP_LOGI(TAG, "Wi-Fi prepared: internal_free=%u largest=%u",
             (unsigned)heap_caps_get_free_size(MALLOC_CAP_INTERNAL),
             (unsigned)heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL));
    return ESP_OK;
}

static esp_err_t stop_station(void)
{
    if (!started) return ESP_OK;
    // Drain the previous attempt's events before accepting a new Wi-Fi job.
    // Otherwise a delayed DISCONNECTED/GOT_IP can affect the next attempt.
    if (!stopping) {
        esp_err_t err = esp_wifi_stop();
        if (err != ESP_OK) return err;
        stopping = true;
    }
    if (!(xEventGroupWaitBits(events, STOPPED, pdFALSE, pdFALSE,
                             pdMS_TO_TICKS(5000)) & STOPPED)) return ESP_ERR_TIMEOUT;
    started = false;
    stopping = false;
    xEventGroupClearBits(events, CONNECTED | DISCONNECTED | STOPPED);
    return ESP_OK;
}

bool ble_wifi_connected(void)
{
    return events && (xEventGroupGetBits(events) & CONNECTED);
}

esp_err_t ble_wifi_connect(const char *ssid, const char *password)
{
    if (!ssid || !ssid[0] || strlen(ssid) > 32 || !password || strlen(password) > 63)
        return ESP_ERR_INVALID_ARG;
    esp_err_t err = ble_wifi_prepare();
    if (err != ESP_OK) return err;
    err = stop_station();
    if (err != ESP_OK) return err;

    wifi_config_t config = {0};
    memcpy(config.sta.ssid, ssid, strlen(ssid));
    memcpy(config.sta.password, password, strlen(password));
    config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    config.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;
    config.sta.sort_method = WIFI_CONNECT_AP_BY_SIGNAL;
    err = esp_wifi_set_config(WIFI_IF_STA, &config);
    if (err != ESP_OK) return err;

    for (int attempt = 0; attempt < 3; ++attempt) {
        xEventGroupClearBits(events, CONNECTED | DISCONNECTED | STOPPED);
        err = esp_wifi_start();
        if (err != ESP_OK) return err;
        started = true;
        err = esp_wifi_connect();
        if (err == ESP_OK) {
            EventBits_t bits = xEventGroupWaitBits(events, CONNECTED | DISCONNECTED,
                                                   pdFALSE, pdFALSE, pdMS_TO_TICKS(10000));
            if (bits & CONNECTED) return ESP_OK;
            err = ESP_ERR_TIMEOUT;
        }
        esp_err_t stop_err = stop_station();
        if (stop_err != ESP_OK) return stop_err;
        if (attempt < 2) vTaskDelay(pdMS_TO_TICKS(500));
    }
    return err;
}