// Wi-Fi station-mode connection manager — see wifi_manager.h for how Wi-Fi
// actually reaches the network on this target (ESP32-C6 co-processor over
// SDIO via esp_wifi_remote/esp_hosted). The esp_wifi_*/esp_netif_* calls
// below are the same ones a Wi-Fi-native chip would use.
//
// Structure follows ESP-IDF's own documented Wi-Fi station example pattern
// (event-driven connect via WIFI_EVENT/IP_EVENT + an event group the rest
// of the app can block on), extended with an exponential-backoff reconnect
// instead of giving up after N attempts — a sensor node left running for
// weeks needs to ride out a router reboot or a home Wi-Fi outage
// indefinitely, not park itself in a permanent failure state.
//
// UNVERIFIED AGAINST REAL HARDWARE — see README.md.

#include <inttypes.h>
#include "wifi_manager.h"
#include "device_config.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "nvs_flash.h"
#include "freertos/event_groups.h"

static const char *TAG = "wifi_manager";

#define WIFI_CONNECTED_BIT BIT0

#define WIFI_RECONNECT_BASE_DELAY_MS 1000
#define WIFI_RECONNECT_MAX_DELAY_MS  30000
#define WIFI_RECONNECT_BACKOFF_CAP_SHIFT 5 // 1000ms << 5 = 32000ms, then clamped to MAX

static EventGroupHandle_t s_wifi_event_group = NULL;
static esp_timer_handle_t s_reconnect_timer = NULL;
static uint32_t s_retry_count = 0;

static void reconnect_timer_cb(void *arg) {
    ESP_LOGI(TAG, "attempting Wi-Fi reconnect");
    esp_err_t err = esp_wifi_connect();
    if (err != ESP_OK) {
        ESP_LOGW(TAG, "esp_wifi_connect() failed: %s", esp_err_to_name(err));
    }
}

static void schedule_reconnect(void) {
    uint32_t shift = s_retry_count < WIFI_RECONNECT_BACKOFF_CAP_SHIFT ? s_retry_count : WIFI_RECONNECT_BACKOFF_CAP_SHIFT;
    uint32_t delay_ms = WIFI_RECONNECT_BASE_DELAY_MS << shift;
    if (delay_ms > WIFI_RECONNECT_MAX_DELAY_MS) {
        delay_ms = WIFI_RECONNECT_MAX_DELAY_MS;
    }
    s_retry_count++;
    ESP_LOGW(TAG, "reconnecting in %" PRIu32 " ms (attempt %" PRIu32 ")", delay_ms, s_retry_count);
    esp_timer_stop(s_reconnect_timer); // no-op if not running; keeps this idempotent
    esp_timer_start_once(s_reconnect_timer, (uint64_t)delay_ms * 1000ULL);
}

static void wifi_event_handler(void *arg, esp_event_base_t event_base,
                                int32_t event_id, void *event_data) {
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        xEventGroupClearBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
        wifi_event_sta_disconnected_t *disc = (wifi_event_sta_disconnected_t *)event_data;
        ESP_LOGW(TAG, "Wi-Fi disconnected (reason %d)", disc ? disc->reason : -1);
        schedule_reconnect();
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *evt = (ip_event_got_ip_t *)event_data;
        ESP_LOGI(TAG, "got IP: " IPSTR, IP2STR(&evt->ip_info.ip));
        s_retry_count = 0;
        xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
    }
}

esp_err_t wifi_manager_start(void) {
    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        err = nvs_flash_init();
    }
    ESP_ERROR_CHECK(err);

    s_wifi_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t init_cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&init_cfg));

    ESP_ERROR_CHECK(esp_event_handler_instance_register(
        WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(
        IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, NULL));

    const esp_timer_create_args_t timer_args = {
        .callback = &reconnect_timer_cb,
        .name = "wifi_reconnect",
    };
    ESP_ERROR_CHECK(esp_timer_create(&timer_args, &s_reconnect_timer));

    wifi_config_t wifi_config = {
        .sta = {
            .ssid = DEVICE_WIFI_SSID,
            .password = DEVICE_WIFI_PASSWORD,
            .threshold.authmode = WIFI_AUTH_WPA2_PSK,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG, "Wi-Fi station starting, SSID \"%s\"", DEVICE_WIFI_SSID);
    return ESP_OK;
}

esp_err_t wifi_manager_wait_connected(TickType_t timeout_ticks) {
    if (s_wifi_event_group == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    EventBits_t bits = xEventGroupWaitBits(
        s_wifi_event_group, WIFI_CONNECTED_BIT,
        pdFALSE /* don't clear on exit */, pdFALSE /* any bit */, timeout_ticks);
    return (bits & WIFI_CONNECTED_BIT) ? ESP_OK : ESP_ERR_TIMEOUT;
}

bool wifi_manager_is_connected(void) {
    if (s_wifi_event_group == NULL) {
        return false;
    }
    return (xEventGroupGetBits(s_wifi_event_group) & WIFI_CONNECTED_BIT) != 0;
}