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

static const char *TAG = "CSI_MGR";

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

// --- 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;
}