/*
 * MTG RFID Companion - Wi-Fi manager (Core 0)
 *
 * Multi-profile station manager with NVS storage, live scanning,
 * and boot auto-connect.
 */

#include <string.h>
#include <stdlib.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "freertos/semphr.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "nvs_flash.h"
#include "nvs.h"
#include "wifi_manager.h"

static const char *TAG = "wifi";

#define WIFI_RETRY_MAX   (UINT8_MAX)
#define NVS_WIFI_NS      "wifi_store"

static EventGroupHandle_t s_wifi_group = NULL;
#define WIFI_CONNECTED_BIT BIT0

static int s_retry_count = 0;
static char s_sta_ip[20] = "No Wi-Fi";
static char s_active_ssid[WIFI_SSID_MAX_LEN + 1] = {0};
static SemaphoreHandle_t s_wifi_mux = NULL;
static SemaphoreHandle_t s_net_lock = NULL;
static bool s_manual_disconnect = false;
static esp_netif_t *s_sta_netif = NULL;
static bool s_power_save_enabled = true;


/* ------------------------------------------------------------------ */
/*  NVS Storage Helpers                                                */
/* ------------------------------------------------------------------ */

int wifi_manager_get_saved_profiles(wifi_profile_t *out_profiles, int max_count)
{
    if (out_profiles == NULL || max_count <= 0) return 0;

    nvs_handle_t h;
    if (nvs_open(NVS_WIFI_NS, NVS_READONLY, &h) != ESP_OK) {
        return 0;
    }

    uint8_t count = 0;
    if (nvs_get_u8(h, "count", &count) != ESP_OK) {
        nvs_close(h);
        return 0;
    }

    int loaded = 0;
    for (uint8_t i = 0; i < count && loaded < max_count; i++) {
        char key_s[16], key_p[16];
        snprintf(key_s, sizeof(key_s), "s_%u", i);
        snprintf(key_p, sizeof(key_p), "p_%u", i);

        size_t s_len = sizeof(out_profiles[loaded].ssid);
        size_t p_len = sizeof(out_profiles[loaded].password);

        if (nvs_get_str(h, key_s, out_profiles[loaded].ssid, &s_len) == ESP_OK &&
            nvs_get_str(h, key_p, out_profiles[loaded].password, &p_len) == ESP_OK) {
            loaded++;
        }
    }

    nvs_close(h);
    return loaded;
}

esp_err_t wifi_manager_save_profile(const char *ssid, const char *password)
{
    if (ssid == NULL || strlen(ssid) == 0) return ESP_ERR_INVALID_ARG;
    if (password == NULL) password = "";

    wifi_profile_t profiles[WIFI_MAX_PROFILES];
    int count = wifi_manager_get_saved_profiles(profiles, WIFI_MAX_PROFILES);

    for (int i = 0; i < count; i++) {
        if (strcmp(profiles[i].ssid, ssid) == 0) {
            nvs_handle_t h;
            esp_err_t err = nvs_open(NVS_WIFI_NS, NVS_READWRITE, &h);
            if (err != ESP_OK) return err;
            char key_p[16];
            snprintf(key_p, sizeof(key_p), "p_%d", i);
            nvs_set_str(h, key_p, password);
            err = nvs_commit(h);
            nvs_close(h);
            ESP_LOGI(TAG, "Updated password for saved profile '%s'", ssid);
            return err;
        }
    }

    if (count >= WIFI_MAX_PROFILES) {
        for (int i = 0; i < WIFI_MAX_PROFILES - 1; i++) {
            profiles[i] = profiles[i + 1];
        }
        count = WIFI_MAX_PROFILES - 1;
    }

    strncpy(profiles[count].ssid, ssid, sizeof(profiles[count].ssid) - 1);
    profiles[count].ssid[sizeof(profiles[count].ssid) - 1] = '\0';
    strncpy(profiles[count].password, password, sizeof(profiles[count].password) - 1);
    profiles[count].password[sizeof(profiles[count].password) - 1] = '\0';
    count++;

    nvs_handle_t h;
    esp_err_t err = nvs_open(NVS_WIFI_NS, NVS_READWRITE, &h);
    if (err != ESP_OK) return err;

    nvs_set_u8(h, "count", (uint8_t)count);
    for (int i = 0; i < count; i++) {
        char key_s[16], key_p[16];
        snprintf(key_s, sizeof(key_s), "s_%d", i);
        snprintf(key_p, sizeof(key_p), "p_%d", i);
        nvs_set_str(h, key_s, profiles[i].ssid);
        nvs_set_str(h, key_p, profiles[i].password);
    }

    err = nvs_commit(h);
    nvs_close(h);
    ESP_LOGI(TAG, "Saved profile '%s' to NVS (total: %d)", ssid, count);
    return err;
}

esp_err_t wifi_manager_delete_profile(const char *ssid)
{
    if (ssid == NULL) return ESP_ERR_INVALID_ARG;

    wifi_profile_t profiles[WIFI_MAX_PROFILES];
    int count = wifi_manager_get_saved_profiles(profiles, WIFI_MAX_PROFILES);

    int del_idx = -1;
    for (int i = 0; i < count; i++) {
        if (strcmp(profiles[i].ssid, ssid) == 0) {
            del_idx = i;
            break;
        }
    }

    if (del_idx == -1) return ESP_ERR_NOT_FOUND;

    for (int i = del_idx; i < count - 1; i++) {
        profiles[i] = profiles[i + 1];
    }
    count--;

    nvs_handle_t h;
    esp_err_t err = nvs_open(NVS_WIFI_NS, NVS_READWRITE, &h);
    if (err != ESP_OK) return err;

    nvs_erase_all(h);
    nvs_set_u8(h, "count", (uint8_t)count);
    for (int i = 0; i < count; i++) {
        char key_s[16], key_p[16];
        snprintf(key_s, sizeof(key_s), "s_%d", i);
        snprintf(key_p, sizeof(key_p), "p_%d", i);
        nvs_set_str(h, key_s, profiles[i].ssid);
        nvs_set_str(h, key_p, profiles[i].password);
    }

    err = nvs_commit(h);
    nvs_close(h);
    ESP_LOGI(TAG, "Deleted profile '%s' from NVS (remaining: %d)", ssid, count);
    return err;
}

/* ------------------------------------------------------------------ */
/*  Event Handler & Diagnostics                                       */
/* ------------------------------------------------------------------ */

static void wifi_event_handler(void *arg, esp_event_base_t base, int32_t id, void *data)
{
    (void)arg;

    if (base == WIFI_EVENT && id == WIFI_EVENT_STA_START) {
        ESP_LOGI(TAG, "Wi-Fi station interface started");
    } else if (base == WIFI_EVENT && id == WIFI_EVENT_STA_DISCONNECTED) {
        strncpy(s_sta_ip, "No Wi-Fi", sizeof(s_sta_ip));
        if (s_wifi_group) {
            xEventGroupClearBits(s_wifi_group, WIFI_CONNECTED_BIT);
        }
        if (!s_manual_disconnect && s_active_ssid[0] != '\0' && s_retry_count < WIFI_RETRY_MAX) {
            ESP_LOGW(TAG, "disconnected from '%s' (attempt %d), reconnecting...", s_active_ssid, s_retry_count);
            esp_wifi_connect();
            s_retry_count++;
        }
    } else if (base == IP_EVENT && id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *event = (ip_event_got_ip_t *)data;
        snprintf(s_sta_ip, sizeof(s_sta_ip), IPSTR, IP2STR(&event->ip_info.ip));
        ESP_LOGI(TAG, "connected to '%s'! Assigned IP: %s", s_active_ssid, s_sta_ip);
        s_retry_count = 0;
        if (s_wifi_group) {
            xEventGroupSetBits(s_wifi_group, WIFI_CONNECTED_BIT);
        }
        if (s_sta_netif != NULL) {
            esp_netif_dns_info_t dns_fallback = {0};
            dns_fallback.ip.type = ESP_IPADDR_TYPE_V4;
            dns_fallback.ip.u_addr.ip4.addr = ESP_IP4TOADDR(8, 8, 8, 8);
            esp_netif_set_dns_info(s_sta_netif, ESP_NETIF_DNS_FALLBACK, &dns_fallback);
        }
        /* Enable modem power saving only if power save is allowed */
        if (s_power_save_enabled) {
            esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
            ESP_LOGI(TAG, "Wi-Fi modem sleep enabled (WIFI_PS_MIN_MODEM: ~30-40 mA saved)");
        } else {
            esp_wifi_set_ps(WIFI_PS_NONE);
            ESP_LOGI(TAG, "Wi-Fi power save disabled (WIFI_PS_NONE)");
        }
    }
}

const char *wifi_manager_get_ip(void)
{
    if (!wifi_manager_connected()) {
        return "No Wi-Fi";
    }
    return s_sta_ip;
}

const char *wifi_manager_get_connected_ssid(void)
{
    if (!wifi_manager_connected() || s_active_ssid[0] == '\0') {
        return "Disconnected";
    }
    return s_active_ssid;
}

int8_t wifi_manager_get_rssi(void)
{
    if (!wifi_manager_connected()) {
        return 0;
    }
    wifi_ap_record_t ap_info;
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        return ap_info.rssi;
    }
    return 0;
}

bool wifi_manager_connected(void)
{
    return s_wifi_group != NULL && (xEventGroupGetBits(s_wifi_group) & WIFI_CONNECTED_BIT) != 0;
}

esp_err_t wifi_manager_wait_connected(TickType_t timeout_ticks)
{
    if (s_wifi_group == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    EventBits_t bits = xEventGroupWaitBits(s_wifi_group, WIFI_CONNECTED_BIT, pdFALSE, pdTRUE, timeout_ticks);
    return (bits & WIFI_CONNECTED_BIT) ? ESP_OK : ESP_ERR_TIMEOUT;
}

bool wifi_manager_sta_configured(void)
{
    wifi_profile_t profiles[1];
    if (wifi_manager_get_saved_profiles(profiles, 1) > 0) {
        return true;
    }
#ifdef CONFIG_MTG_WIFI_SSID
    const char *ssid = CONFIG_MTG_WIFI_SSID;
    return strlen(ssid) > 0 && strstr(ssid, "YOUR_") == NULL;
#else
    return false;
#endif
}

/* ------------------------------------------------------------------ */
/*  Scanning & Connection APIs                                        */
/* ------------------------------------------------------------------ */

esp_err_t wifi_manager_scan_networks(wifi_ap_record_t **out_records, uint16_t *out_count)
{
    if (out_records == NULL || out_count == NULL) return ESP_ERR_INVALID_ARG;
    *out_records = NULL;
    *out_count = 0;

    if (s_wifi_mux) xSemaphoreTake(s_wifi_mux, portMAX_DELAY);

    wifi_scan_config_t scan_cfg = {
        .show_hidden = true,
    };

    ESP_LOGI(TAG, "Scanning for Wi-Fi networks...");
    esp_err_t err = esp_wifi_scan_start(&scan_cfg, true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_wifi_scan_start failed: %s", esp_err_to_name(err));
        if (s_wifi_mux) xSemaphoreGive(s_wifi_mux);
        return err;
    }

    uint16_t ap_count = 0;
    ESP_ERROR_CHECK(esp_wifi_scan_get_ap_num(&ap_count));

    if (ap_count == 0) {
        if (s_wifi_mux) xSemaphoreGive(s_wifi_mux);
        return ESP_OK;
    }

    wifi_ap_record_t *records = malloc(sizeof(wifi_ap_record_t) * ap_count);
    if (records == NULL) {
        if (s_wifi_mux) xSemaphoreGive(s_wifi_mux);
        return ESP_ERR_NO_MEM;
    }

    esp_err_t ret = esp_wifi_scan_get_ap_records(&ap_count, records);
    if (ret == ESP_OK) {
        *out_records = records;
        *out_count = ap_count;
        ESP_LOGI(TAG, "Scan complete: found %u access points", ap_count);
    } else {
        free(records);
    }

    if (s_wifi_mux) xSemaphoreGive(s_wifi_mux);
    return ret;
}

esp_err_t wifi_manager_connect_new(const char *ssid, const char *password, bool save_to_nvs)
{
    if (ssid == NULL || strlen(ssid) == 0) {
        return ESP_ERR_INVALID_ARG;
    }

    ESP_LOGI(TAG, "Connecting to new network: '%s'", ssid);

    wifi_config_t wifi_cfg = {0};
    strncpy((char *)wifi_cfg.sta.ssid, ssid, sizeof(wifi_cfg.sta.ssid) - 1);
    if (password && strlen(password) > 0) {
        strncpy((char *)wifi_cfg.sta.password, password, sizeof(wifi_cfg.sta.password) - 1);
        wifi_cfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    } else {
        wifi_cfg.sta.threshold.authmode = WIFI_AUTH_OPEN;
    }

    strncpy(s_active_ssid, ssid, sizeof(s_active_ssid) - 1);
    s_active_ssid[sizeof(s_active_ssid) - 1] = '\0';
    s_manual_disconnect = true;
    esp_wifi_disconnect();
    s_manual_disconnect = false;
    esp_err_t set_err = esp_wifi_set_config(WIFI_IF_STA, &wifi_cfg);
    if (set_err != ESP_OK) {
        ESP_LOGE(TAG, "esp_wifi_set_config failed: %s", esp_err_to_name(set_err));
        return set_err;
    }
    esp_err_t ret = esp_wifi_connect();

    if (save_to_nvs) {
        wifi_manager_save_profile(ssid, password);
    }

    return ret;
}

/* ------------------------------------------------------------------ */
/*  Subsystem Init & Boot Auto-Connect                                */
/* ------------------------------------------------------------------ */

esp_err_t wifi_manager_init(void)
{
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    if (ret != ESP_OK) {
        return ret;
    }

    s_wifi_mux = xSemaphoreCreateMutex();
    s_net_lock = xSemaphoreCreateMutex();
    s_wifi_group = xEventGroupCreate();
    if (s_wifi_group == NULL || s_net_lock == NULL) {
        return ESP_ERR_NO_MEM;
    }


    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    s_sta_netif = esp_netif_create_default_wifi_sta();
    if (s_sta_netif != NULL) {
        esp_netif_set_hostname(s_sta_netif, "mtg-companion");
    }

    esp_netif_t *ap_netif = esp_netif_create_default_wifi_ap();
    if (ap_netif != NULL) {
        esp_netif_set_hostname(ap_netif, "mtg-companion");
        esp_netif_ip_info_t ip_info = {
            .ip = { .addr = ESP_IP4TOADDR(192, 168, 4, 1) },
            .netmask = { .addr = ESP_IP4TOADDR(255, 255, 255, 0) },
            .gw = { .addr = ESP_IP4TOADDR(192, 168, 4, 1) },
        };
        esp_netif_dhcps_stop(ap_netif);
        esp_netif_set_ip_info(ap_netif, &ip_info);
        esp_netif_dhcps_start(ap_netif);
    }

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_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));

    return ESP_OK;
}

esp_err_t wifi_manager_start(void)
{
    wifi_profile_t saved_profiles[WIFI_MAX_PROFILES];
    int profile_count = wifi_manager_get_saved_profiles(saved_profiles, WIFI_MAX_PROFILES);

    const char *target_ssid = NULL;
    const char *target_pass = NULL;

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_APSTA));

    wifi_config_t ap_cfg = {
        .ap = {
            .ssid_len = sizeof("MTG-Companion") - 1,
            .max_connection = 4,
            .authmode = WIFI_AUTH_OPEN,
        },
    };
    strncpy((char *)ap_cfg.ap.ssid, "MTG-Companion", sizeof(ap_cfg.ap.ssid) - 1);
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &ap_cfg));

    ESP_ERROR_CHECK(esp_wifi_start());

    if (profile_count > 0) {
        ESP_LOGI(TAG, "Found %d saved Wi-Fi profiles in NVS. Scanning for best match...", profile_count);

        wifi_scan_config_t scan_cfg = { .show_hidden = true };
        esp_wifi_scan_start(&scan_cfg, true);

        uint16_t ap_count = 0;
        esp_wifi_scan_get_ap_num(&ap_count);

        int best_profile_idx = -1;
        int best_rssi = -127;

        if (ap_count > 0) {
            wifi_ap_record_t *ap_list = malloc(sizeof(wifi_ap_record_t) * ap_count);
            if (ap_list && esp_wifi_scan_get_ap_records(&ap_count, ap_list) == ESP_OK) {
                for (uint16_t a = 0; a < ap_count; a++) {
                    for (int p = 0; p < profile_count; p++) {
                        if (strcmp((char *)ap_list[a].ssid, saved_profiles[p].ssid) == 0) {
                            if (ap_list[a].rssi > best_rssi) {
                                best_rssi = ap_list[a].rssi;
                                best_profile_idx = p;
                            }
                        }
                    }
                }
            }
            free(ap_list);
        }

        if (best_profile_idx >= 0) {
            ESP_LOGI(TAG, "Matched in-range network '%s' (RSSI %d dBm)",
                     saved_profiles[best_profile_idx].ssid, best_rssi);
            target_ssid = saved_profiles[best_profile_idx].ssid;
            target_pass = saved_profiles[best_profile_idx].password;
        } else {
            ESP_LOGW(TAG, "No saved networks currently detected in scan. Trying first saved profile '%s'...",
                     saved_profiles[0].ssid);
            target_ssid = saved_profiles[0].ssid;
            target_pass = saved_profiles[0].password;
        }
    } else {
        /* Fall back to compile-time Kconfig defaults */
#ifdef CONFIG_MTG_WIFI_SSID
        if (wifi_manager_sta_configured()) {
            ESP_LOGI(TAG, "No NVS profiles. Using Kconfig default '%s'...", CONFIG_MTG_WIFI_SSID);
            target_ssid = CONFIG_MTG_WIFI_SSID;
            target_pass = CONFIG_MTG_WIFI_PASS;
        }
#endif
    }

    if (target_ssid != NULL && target_ssid[0] != '\0') {
        wifi_config_t wifi_cfg = {0};
        strncpy((char *)wifi_cfg.sta.ssid, target_ssid, sizeof(wifi_cfg.sta.ssid) - 1);
        if (target_pass && target_pass[0] != '\0') {
            strncpy((char *)wifi_cfg.sta.password, target_pass, sizeof(wifi_cfg.sta.password) - 1);
            wifi_cfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
        } else {
            wifi_cfg.sta.threshold.authmode = WIFI_AUTH_OPEN;
        }

        strncpy(s_active_ssid, target_ssid, sizeof(s_active_ssid) - 1);
        s_active_ssid[sizeof(s_active_ssid) - 1] = '\0';
        s_manual_disconnect = true;
        esp_wifi_disconnect();
        s_manual_disconnect = false;
        esp_err_t err = esp_wifi_set_config(WIFI_IF_STA, &wifi_cfg);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_wifi_set_config failed: %s", esp_err_to_name(err));
        } else {
            ESP_LOGI(TAG, "Connecting to target '%s'...", target_ssid);
            esp_wifi_connect();
        }
    } else {
        ESP_LOGI(TAG, "No Wi-Fi networks configured yet. Device running in SoftAP setup mode.");
    }

    return ESP_OK;
}

esp_err_t wifi_manager_acquire_net_lock(TickType_t timeout_ticks)
{
    if (!s_net_lock) {
        s_net_lock = xSemaphoreCreateMutex();
    }
    if (!s_net_lock) {
        return ESP_ERR_NO_MEM;
    }
    BaseType_t taken = xSemaphoreTake(s_net_lock, timeout_ticks);
    if (taken == pdTRUE) {
        /* Full RF performance for active TLS data transfers / OTA */
        esp_wifi_set_ps(WIFI_PS_NONE);
        return ESP_OK;
    }
    return ESP_ERR_TIMEOUT;
}

void wifi_manager_release_net_lock(void)
{
    if (s_net_lock) {
        /* Revert to low-power modem sleep when idle */
        if (s_power_save_enabled && wifi_manager_connected()) {
            esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
        }
        xSemaphoreGive(s_net_lock);
    }
}

esp_err_t wifi_manager_set_power_save(bool enable)
{
    s_power_save_enabled = enable;
    wifi_ps_type_t ps_type = enable ? WIFI_PS_MIN_MODEM : WIFI_PS_NONE;
    esp_err_t err = esp_wifi_set_ps(ps_type);
    ESP_LOGI(TAG, "Wi-Fi power save %s (%s)",
             enable ? "enabled" : "disabled",
             enable ? "WIFI_PS_MIN_MODEM" : "WIFI_PS_NONE");
    return err;
}