#include "csi_manager.h"
#include "csi_log.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "nvs.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

static const char *TAG = "CSI_MGR";
static const char *NVS_NAMESPACE = "csi_config";
static const char *NVS_KEY_ENABLE = "csi_enable";

static bool s_csi_enabled = false;
static uint32_t s_csi_packet_count = 0;

// --- NVS Functions ---

esp_err_t csi_mgr_save_enable_state(bool enable) {
    nvs_handle_t handle;
    esp_err_t err;

    err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to open NVS: %s", esp_err_to_name(err));
        return err;
    }

    err = nvs_set_u8(handle, NVS_KEY_ENABLE, enable ? 1 : 0);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to write CSI enable state: %s", esp_err_to_name(err));
        nvs_close(handle);
        return err;
    }

    err = nvs_commit(handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to commit NVS: %s", esp_err_to_name(err));
    } else {
        ESP_LOGI(TAG, "CSI enable state saved: %s", enable ? "ENABLED" : "DISABLED");
    }

    nvs_close(handle);
    return err;
}

esp_err_t csi_mgr_load_enable_state(bool *enable) {
    if (!enable) {
        return ESP_ERR_INVALID_ARG;
    }

    nvs_handle_t handle;
    esp_err_t err;

    err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        if (err == ESP_ERR_NVS_NOT_FOUND) {
            ESP_LOGW(TAG, "CSI config namespace not found - using default (disabled)");
            *enable = false;
            return ESP_ERR_NVS_NOT_FOUND;
        }
        ESP_LOGE(TAG, "Failed to open NVS: %s", esp_err_to_name(err));
        *enable = false;
        return err;
    }

    uint8_t value = 0;
    err = nvs_get_u8(handle, NVS_KEY_ENABLE, &value);
    nvs_close(handle);

    if (err == ESP_ERR_NVS_NOT_FOUND) {
        ESP_LOGI(TAG, "CSI enable not configured - using default (disabled)");
        *enable = false;
        return ESP_ERR_NVS_NOT_FOUND;
    } else if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to read CSI enable state: %s", esp_err_to_name(err));
        *enable = false;
        return err;
    }

    *enable = (value != 0);
    ESP_LOGI(TAG, "CSI enable loaded from NVS: %s", *enable ? "ENABLED" : "DISABLED");
    return ESP_OK;
}

bool csi_mgr_should_enable(void) {
    bool enable = false;
    csi_mgr_load_enable_state(&enable);
    return enable;
}

// --- Callback ---
static void csi_cb(void *ctx, wifi_csi_info_t *info) {
    if (!info || !s_csi_enabled) return;

    csi_log_append_record(info);
    s_csi_packet_count++;

    if ((s_csi_packet_count % 100) == 0) {
        ESP_LOGI(TAG, "Captured %lu CSI packets", (unsigned long)s_csi_packet_count);
    }
}

// --- Tasks ---

static void enable_task(void *arg) {
    // Wait for connection to stabilize
    vTaskDelay(pdMS_TO_TICKS(2000));

    if (s_csi_enabled) {
        vTaskDelete(NULL);
        return;
    }

    wifi_csi_config_t csi_cfg = { 0 };

#if defined(CONFIG_IDF_TARGET_ESP32S3) || defined(CONFIG_IDF_TARGET_ESP32) || defined(CONFIG_IDF_TARGET_ESP32S2)
    // Detailed config for Xtensa chips
    csi_cfg.lltf_en = true;
    csi_cfg.htltf_en = true;
    csi_cfg.stbc_htltf2_en = true;
    csi_cfg.ltf_merge_en = true;
    csi_cfg.channel_filter_en = true;
    csi_cfg.manu_scale = false;
    csi_cfg.shift = false;
#endif

    ESP_LOGI(TAG, "Configuring CSI...");
    if (esp_wifi_set_csi_config(&csi_cfg) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to set CSI config");
        vTaskDelete(NULL);
        return;
    }

    if (esp_wifi_set_csi_rx_cb(csi_cb, NULL) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to set CSI callback");
        vTaskDelete(NULL);
        return;
    }

    if (esp_wifi_set_csi(true) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to enable CSI");
        vTaskDelete(NULL);
        return;
    }

    s_csi_enabled = true;
    ESP_LOGI(TAG, "CSI enabled!");
    vTaskDelete(NULL);
}

static void dump_task(void *arg) {
    // Wait before dumping (per your original logic)
    vTaskDelay(pdMS_TO_TICKS(20000));

    ESP_LOGI(TAG, "Dumping CSI data...");
    csi_log_dump_over_uart();
    ESP_LOGI(TAG, "CSI dump complete");
    vTaskDelete(NULL);
}

// --- API ---

void csi_mgr_init(void) {
    s_csi_enabled = false;
    s_csi_packet_count = 0;
}

void csi_mgr_enable_async(void) {
    if (s_csi_enabled) return;
    xTaskCreate(enable_task, "csi_enable", 4096, NULL, 5, NULL);
}

void csi_mgr_disable(void) {
    if (!s_csi_enabled) return;

    ESP_LOGI(TAG, "Disabling CSI...");
    esp_wifi_set_csi(false);
    s_csi_enabled = false;
    ESP_LOGI(TAG, "CSI disabled");
}

void csi_mgr_schedule_dump(void) {
    xTaskCreate(dump_task, "csi_dump", 4096, NULL, 5, NULL);
}

bool csi_mgr_is_enabled(void) {
    return s_csi_enabled;
}

uint32_t csi_mgr_get_packet_count(void) {
    return s_csi_packet_count;
}