#include "csi_log.h"

#include <string.h>
#include <stdio.h>
#include <inttypes.h>

#include "esp_partition.h"
#include "esp_log.h"

static const char *TAG = "storage";

typedef struct __attribute__((packed)) {
    uint32_t seq;
    int32_t  ts_us;
    int8_t   rssi;
    uint8_t  reserved;
    uint16_t len;     // CSI payload length in bytes
} csi_record_hdr_t;

static const esp_partition_t *s_csi_part = NULL;
static size_t s_csi_offset = 0;   // bytes written this run
static uint32_t s_csi_seq = 0;

esp_err_t csi_log_init(void) {
    s_csi_part = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, \
                                          ESP_PARTITION_SUBTYPE_DATA_SPIFFS, "storage");

    if (!s_csi_part) {
        ESP_LOGE(TAG, "csi_log partition not found");
        return ESP_ERR_NOT_FOUND;
    }

    ESP_LOGI(TAG, "csi_log size=%u bytes, addr=0x%08x",
             (unsigned)s_csi_part->size,
             (unsigned)s_csi_part->address);

    // Erase once per experiment
    esp_err_t err = esp_partition_erase_range(
        s_csi_part,
        0,
        s_csi_part->size);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "erase failed: %s", esp_err_to_name(err));
        s_csi_part = NULL;
        return err;
    }

    s_csi_offset = 0;
    s_csi_seq = 0;
    return ESP_OK;
}

void csi_log_append_record(const wifi_csi_info_t *info)
{
    if (!s_csi_part || !info) {
        return;
    }

    // Basic sanity bounds
    if (info->len == 0 || info->len > 512) {
        return;
    }

    csi_record_hdr_t hdr = {
        .seq      = s_csi_seq++,
        .ts_us    = (int32_t)info->rx_ctrl.timestamp,
        .rssi     = info->rx_ctrl.rssi,
        .reserved = 0,
        .len      = info->len,
    };

    uint16_t rec_bytes = sizeof(hdr) + info->len;
    size_t needed = s_csi_offset + 2 /*len field*/ + rec_bytes;

    if (needed > s_csi_part->size) {
        static bool warned = false;
        if (!warned) {
            ESP_LOGW(TAG, "csi_log full, dropping subsequent records");
            warned = true;
        }
        return;
    }

    esp_err_t err;

    // Write record length (header+payload)
    err = esp_partition_write(s_csi_part, s_csi_offset,
                              &rec_bytes, sizeof(rec_bytes));
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "write rec_len failed: %s", esp_err_to_name(err));
        return;
    }
    s_csi_offset += sizeof(rec_bytes);

    // Write header
    err = esp_partition_write(s_csi_part, s_csi_offset,
                              &hdr, sizeof(hdr));
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "write hdr failed: %s", esp_err_to_name(err));
        return;
    }
    s_csi_offset += sizeof(hdr);

    // Write payload
    err = esp_partition_write(s_csi_part, s_csi_offset,
                              info->buf, info->len);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "write payload failed: %s", esp_err_to_name(err));
        return;
    }
    s_csi_offset += info->len;
}

void csi_log_dump_over_uart(void)
{
    if (!s_csi_part) {
        printf("CSI_DUMP_ERROR no_partition\n");
        return;
    }

    printf("CSI_DUMP_BEGIN %u\n", (unsigned)s_csi_offset);

    size_t pos = 0;
    while (pos + 2 <= s_csi_offset) {
        uint16_t rec_bytes = 0;
        esp_err_t err = esp_partition_read(
            s_csi_part, pos, &rec_bytes, sizeof(rec_bytes));
        if (err != ESP_OK) {
            printf("CSI_DUMP_ERROR read_len %s\n", esp_err_to_name(err));
            break;
        }
        pos += sizeof(rec_bytes);

        if (rec_bytes < sizeof(csi_record_hdr_t) ||
            pos + rec_bytes > s_csi_offset) {
            break;
        }

        csi_record_hdr_t hdr;
        err = esp_partition_read(
            s_csi_part, pos, &hdr, sizeof(hdr));
        if (err != ESP_OK) {
            printf("CSI_DUMP_ERROR read_hdr %s\n", esp_err_to_name(err));
            break;
        }
        pos += sizeof(hdr);

        uint16_t csi_len = hdr.len;
        if (csi_len > rec_bytes - sizeof(hdr)) {
            break;
        }

        static uint8_t buf[512];
        if (csi_len > sizeof(buf)) {
            csi_len = sizeof(buf);
        }

        err = esp_partition_read(
            s_csi_part, pos, buf, csi_len);
        if (err != ESP_OK) {
            printf("CSI_DUMP_ERROR read_csi %s\n", esp_err_to_name(err));
            break;
        }
        pos += (rec_bytes - sizeof(hdr)); // skip to record boundary

        // Text record line
        printf("R %" PRIu32 " %" PRId32 " %d %u ",
               hdr.seq, hdr.ts_us, hdr.rssi, (unsigned)csi_len);
        for (int i = 0; i < csi_len; i++) {
            printf("%02X", buf[i]);
        }
        printf("\n");
    }

    printf("CSI_DUMP_END\n");
}