/**
 * app_main.c - C-ITS Sniffer ESP32-C5
 *
 * IEEE 802.15.4 Promiscuous Mode Sniffer
 * Ringpuffer für JTAG-Auslesung von rohen Paketen
 * Konsolencommands: status, read_pkts [n], chan <ch>, packets <n>
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <time.h>
#include <inttypes.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include "sdkconfig.h"
#include "esp_err.h"
#include "esp_log.h"
#include "esp_log_write.h"
#include "esp_timer.h"
#include "esp_console.h"
#include "esp_vfs.h"
#include "driver/usb_serial_jtag.h"
#include "driver/usb_serial_jtag_vfs.h"
#include "nvs_flash.h"
#include "esp_vfs_fat.h"
#include "wear_levelling.h"
#include "esp_ieee802154.h"
#include "esp_ieee802154_types.h"
#include "driver/gpio.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

static const char *TAG = "CITS_SNIFFER";
static const char *TAG_CMD = "CMD";
#define LED_PIN GPIO_NUM_28
#define SNIFFER_CHANNEL  11
#define CSV_INTERVAL_SEC 3600

// Ringpuffer für rohe Pakete (für JTAG-Direct-Access)
#define PKT_RING_BUF_SIZE 128
#define PKT_MAX_SIZE 160
typedef struct {
    uint8_t data[PKT_MAX_SIZE];
    uint16_t len;
    int8_t rssi;
    uint8_t lqi;
    uint32_t ts_sec;
    uint32_t ts_usec;
} raw_pkt_t;

static raw_pkt_t pkt_ring[PKT_RING_BUF_SIZE];
static volatile uint32_t pkt_ring_write = 0;
static volatile uint32_t pkt_ring_total = 0;
static bool pkt_ring_full = false;

static FILE *csv_file = NULL;
static char csv_filename[64];
static time_t last_csv_time = 0;
static uint32_t packet_count = 0;
static wl_handle_t wl_handle = WL_INVALID_HANDLE;

static volatile bool sniffer_active = false;
static volatile int8_t last_rssi = -127;

// === CSV / Storage ===

static void write_csv_header(FILE *fp)
{
    fprintf(fp, "timestamp_sec,timestamp_usec,length,sender_mac,frame_type,channel,rssi,lqi\n");
}

static void write_csv_row(FILE *fp, uint32_t ts_sec, uint32_t ts_usec, uint16_t len,
                           uint8_t *src_mac, uint8_t ftype, uint8_t ch, int16_t rssi, uint8_t lqi)
{
    char mac_str[24];
    snprintf(mac_str, sizeof(mac_str), "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x",
             src_mac[7], src_mac[6], src_mac[5], src_mac[4], src_mac[3], src_mac[2], src_mac[1], src_mac[0]);
    fprintf(fp, "%" PRIu32 ",%" PRIu32 ",%" PRIu16 ",%s,%u,%u,%d,%u\n",
            ts_sec, ts_usec, len, mac_str, ftype, ch, rssi, lqi);
}

static esp_err_t create_csv_file(void)
{
    time_t now;
    struct tm *tm_info;
    time(&now);
    tm_info = localtime(&now);
    snprintf(csv_filename, sizeof(csv_filename),
             "/storage/cits_%04d%02d%02d_%02d%02d%02d.csv",
             tm_info->tm_year + 1900, tm_info->tm_mon + 1, tm_info->tm_mday,
             tm_info->tm_hour, tm_info->tm_min, tm_info->tm_sec);
    if (csv_file != NULL) fclose(csv_file);
    csv_file = fopen(csv_filename, "w");
    if (csv_file == NULL) {
        ESP_LOGE(TAG, "CSV create failed: %s", csv_filename);
        return ESP_FAIL;
    }
    write_csv_header(csv_file);
    fflush(csv_file);
    ESP_LOGI(TAG, "CSV: %s", csv_filename);
    return ESP_OK;
}

static void store_pkt_raw(const uint8_t *data, uint16_t len, int8_t rssi, uint8_t lqi,
                           uint32_t ts_sec, uint32_t ts_usec)
{
    uint32_t idx = pkt_ring_write;
    pkt_ring[idx].len = (len > PKT_MAX_SIZE) ? PKT_MAX_SIZE : len;
    memcpy(pkt_ring[idx].data, data, pkt_ring[idx].len);
    pkt_ring[idx].rssi = rssi;
    pkt_ring[idx].lqi = lqi;
    pkt_ring[idx].ts_sec = ts_sec;
    pkt_ring[idx].ts_usec = ts_usec;
    pkt_ring_write = (pkt_ring_write + 1) % PKT_RING_BUF_SIZE;
    pkt_ring_total++;
    if (pkt_ring_total > PKT_RING_BUF_SIZE) pkt_ring_full = true;
}

// === 802.15.4 Sniffer Callback ===

static void sniffer_rx_cb(uint8_t *data, esp_ieee802154_frame_info_t *info)
{
    if (!data || !info) return;
    if (info->rssi < -100) return;
    if (!sniffer_active) return;

    packet_count++;
    last_rssi = info->rssi;

    uint16_t frame_ctrl = data[0] | (data[1] << 8);
    uint8_t ftype = frame_ctrl & 0x07;
    uint8_t addr_mode_src = (frame_ctrl >> 12) & 0x03;

    uint8_t src_mac[8] = {0};
    if (addr_mode_src == 1) {
        src_mac[7] = data[2]; src_mac[6] = data[3];
    } else if (addr_mode_src == 2) {
        memcpy(src_mac, &data[2], 8);
    }

    uint32_t ts_sec = (uint32_t)(info->timestamp / 1000000);
    uint32_t ts_usec = (uint32_t)(info->timestamp / 1000) % 1000000;

    // Raw-Paket im Ringpuffer speichern
    store_pkt_raw(data, (uint16_t)(info->timestamp >> 8), info->rssi, info->lqi, ts_sec, ts_usec);

    if (csv_file) {
        write_csv_row(csv_file, ts_sec, ts_usec, 0, src_mac, ftype,
                      info->channel, info->rssi, info->lqi);
        fflush(csv_file);
    }

    if (packet_count % 5 == 0) {
        char mac_str[32];
        snprintf(mac_str, sizeof(mac_str), "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x",
                 src_mac[7], src_mac[6], src_mac[5], src_mac[4], src_mac[3], src_mac[2],
                 src_mac[1], src_mac[0]);
        ESP_LOGI(TAG, "[Pkt %lu] RSSI=%d | LQI=%u | Ch=%u | Src=%s",
                 (unsigned long)packet_count, info->rssi, info->lqi,
                 info->channel, mac_str);
    }
}

static esp_err_t init_sniffer(void)
{
    ESP_LOGI(TAG, "Starte 802.15.4 Stack...");
    esp_err_t ret = esp_ieee802154_enable();
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "esp_ieee802154_enable failed: %s", esp_err_to_name(ret));
        return ret;
    }
    ESP_LOGI(TAG, "802.15.4 stack enabled");

    ret = esp_ieee802154_set_channel(SNIFFER_CHANNEL);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "channel set failed: %s", esp_err_to_name(ret));
        return ret;
    }
    ESP_LOGI(TAG, "Channel: %d", SNIFFER_CHANNEL);

    ret = esp_ieee802154_set_promiscuous(true);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "promiscuous failed: %s", esp_err_to_name(ret));
        return ret;
    }
    ESP_LOGI(TAG, "Promiscuous mode enabled");

    esp_ieee802154_event_cb_list_t cb_list = { .rx_done_cb = sniffer_rx_cb };
    ret = esp_ieee802154_event_callback_list_register(cb_list);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "callback register failed: %s", esp_err_to_name(ret));
        return ret;
    }
    ESP_LOGI(TAG, "Rx callback registered");

    sniffer_active = true;
    return ESP_OK;
}

static esp_err_t init_storage(void)
{
    static wl_handle_t wl;
    const esp_vfs_fat_mount_config_t mcfg = {
        .max_files = 4, .format_if_mount_failed = true, .allocation_unit_size = 16 * 1024,
    };
    esp_err_t ret = esp_vfs_fat_spiflash_mount_rw_wl("/storage", "storage", &mcfg, &wl);
    if (ret == ESP_OK) {
        wl_handle = wl;
        ESP_LOGI(TAG, "FATFS storage mounted");
    }
    return ret;
}

// === Konsolencommands ===

static int cmd_status(int argc, char **argv)
{
    (void)argc; (void)argv;
    ESP_LOGI(TAG_CMD, "== Status ==");
    ESP_LOGI(TAG_CMD, "Packets: %lu", (unsigned long)packet_count);
    ESP_LOGI(TAG_CMD, "CSV: %s", csv_filename);
    ESP_LOGI(TAG_CMD, "Channel: %d", SNIFFER_CHANNEL);
    ESP_LOGI(TAG_CMD, "Ring buffer: %lu/%d total, write_idx=%lu",
             (unsigned long)pkt_ring_total, PKT_RING_BUF_SIZE,
             (unsigned long)pkt_ring_write);
    ESP_LOGI(TAG_CMD, "Last RSSI: %d", last_rssi);
    ESP_LOGI(TAG_CMD, "Sniffer active: %s", sniffer_active ? "YES" : "NO");
    return 0;
}

static int cmd_packets(int argc, char **argv)
{
    if (argc < 2) { ESP_LOGW(TAG_CMD, "Usage: packets <count>"); return 1; }
    int count = atoi(argv[1]);
    ESP_LOGI(TAG_CMD, "Warte auf %d Pakete...", count);
    uint32_t start = packet_count;
    while (packet_count - start < count) vTaskDelay(pdMS_TO_TICKS(1000));
    ESP_LOGI(TAG_CMD, "%d Pakete empfangen!", count);
    return 0;
}

static int cmd_read_pkts(int argc, char **argv)
{
    int count = 10;
    if (argc > 0) {
        count = atoi(argv[0]);
        if (count < 1) count = 1;
        if (count > PKT_RING_BUF_SIZE) count = PKT_RING_BUF_SIZE;
    }

    uint32_t total = pkt_ring_total;
    if (total == 0) {
        ESP_LOGI(TAG_CMD, "Keine Pakete empfangen.");
        return 0;
    }

    ESP_LOGI(TAG_CMD, "=== Ring Buffer Read (%d von %lu) ===", count, (unsigned long)total);
    uint32_t idx = (pkt_ring_write == 0) ? PKT_RING_BUF_SIZE - 1 : pkt_ring_write - 1;

    for (int i = 0; i < count; i++) {
        uint32_t read_total = pkt_ring_total;
        if (pkt_ring_full && (read_total - (pkt_ring_write == 0 ? PKT_RING_BUF_SIZE : pkt_ring_write)) <= i) {
            break;
        }
        if (!pkt_ring_full && i >= pkt_ring_total) break;

        raw_pkt_t *pkt = &pkt_ring[idx];
        char hex_str[72];
        hex_str[0] = 0;
        for (int j = 0; j < pkt->len && j < 36; j++) {
            snprintf(hex_str + strlen(hex_str), sizeof(hex_str) - strlen(hex_str),
                     "%02x ", pkt->data[j]);
        }
        ESP_LOGI(TAG_CMD, "  [%d] len=%u RSSI=%d LQI=%u | %s",
                 i, pkt->len, pkt->rssi, pkt->lqi, hex_str);
        idx = (idx == 0) ? PKT_RING_BUF_SIZE - 1 : idx - 1;
    }
    return 0;
}

static int cmd_chan(int argc, char **argv)
{
    if (argc < 1) {
        ESP_LOGW(TAG_CMD, "Usage: chan <channel>");
        return 1;
    }
    int ch = atoi(argv[0]);
    if (ch < 11 || ch > 26) {
        ESP_LOGW(TAG_CMD, "Channel muss 11-26 sein");
        return 1;
    }
    esp_err_t ret = esp_ieee802154_set_channel(ch);
    if (ret == ESP_OK) {
        ESP_LOGI(TAG_CMD, "Channel auf %d geändert", ch);
    } else {
        ESP_LOGE(TAG_CMD, "Channel change failed: %s", esp_err_to_name(ret));
    }
    return 0;
}

static void register_console_commands(void)
{
    esp_console_cmd_t cmd_status_def = { .command = "status", .help = "Show sniffer status", .func = cmd_status };
    esp_console_cmd_t cmd_packets_def = { .command = "packets", .help = "Wait for N packets", .func = cmd_packets };
    esp_console_cmd_t cmd_read_def = { .command = "read_pkts", .help = "Read raw packets from ring buffer", .func = cmd_read_pkts };
    esp_console_cmd_t cmd_ch_def = { .command = "chan", .help = "Change 802.15.4 channel", .func = cmd_chan };
    esp_console_register_help_command();
    esp_console_cmd_register(&cmd_status_def);
    esp_console_cmd_register(&cmd_packets_def);
    esp_console_cmd_register(&cmd_read_def);
    esp_console_cmd_register(&cmd_ch_def);
    ESP_LOGI(TAG, "Console commands registriert");
}

// === JTAG vprintf-Wrapper ===
// Sendet Logs über JTAG-Verbindung zum Host
static volatile int jtag_connected = 0;
static TickType_t jtag_last_connect_tick = 0;

static int jtag_vprintf(const char *fmt, va_list args)
{
    char buf[256];
    int len = vsnprintf(buf, sizeof(buf), fmt, args);
    if (len > 0 && len < 256) {
        // Versuche non-blocking JTAG write
        TickType_t now = xTaskGetTickCount();
        if (now - jtag_last_connect_tick < pdMS_TO_TICKS(1000)) {
            // JTAG write with short timeout (100ms)
            int wrote = usb_serial_jtag_write_bytes(buf, len, pdMS_TO_TICKS(100));
            if (wrote <= 0) {
                jtag_connected = 0;
                jtag_last_connect_tick = now;
            }
        }
    }
    return len;
}

void app_main(void)
{
    ESP_LOGI(TAG, "== C-ITS SNIFFER ESP32-C5 ==");
    ESP_LOGI(TAG, "Build: %s %s", __DATE__, __TIME__);

    // LED initialisieren
    gpio_reset_pin(LED_PIN);
    gpio_set_direction(LED_PIN, GPIO_MODE_OUTPUT);
    gpio_set_level(LED_PIN, 0);

    // NVS initialisieren
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        nvs_flash_erase();
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    // === USB-Serial/JTAG initialisieren ===
    ESP_LOGI(TAG, "Init USB-Serial/JTAG...");
    usb_serial_jtag_driver_config_t jtag_config = {0};
    ret = usb_serial_jtag_driver_install(&jtag_config);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "usb_serial_jtag_driver_install failed: %s", esp_err_to_name(ret));
        return;
    }
    ret = usb_serial_jtag_vfs_register();
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "usb_serial_jtag_vfs_register failed: %s", esp_err_to_name(ret));
        return;
    }

    // Non-blocking JTAG write aktivieren
    usb_serial_jtag_vfs_use_nonblocking();
    ESP_LOGI(TAG, "Non-blocking JTAG aktiv");

    // Log-Ausgabe auf JTAG umleiten
    esp_log_set_vprintf((vprintf_like_t)jtag_vprintf);

    ESP_LOGI(TAG, "USB-Serial/JTAG initialisiert");

    // Storage initialisieren
    ESP_LOGI(TAG, "Init storage...");
    ESP_ERROR_CHECK(init_storage());

    // Console commands registrieren
    register_console_commands();

    // 802.15.4 Sniffer starten
    ret = init_sniffer();
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "802.15.4 init fehlgeschlagen: %s", esp_err_to_name(ret));
        ESP_LOGI(TAG, "Sniffer deaktiviert - nur Console verfügbar");
    }

    // CSV-Datei erstellen
    if (ret == ESP_OK) {
        ESP_ERROR_CHECK(create_csv_file());
        last_csv_time = time(NULL);
    }

    ESP_LOGI(TAG, "== SNIFFER LÄUFT ==");
    ESP_LOGI(TAG, "Channel %d | Warte auf Pakete...", SNIFFER_CHANNEL);
    ESP_LOGI(TAG, "Console: sende 'status', 'read_pkts [n]', 'chan <ch>', 'packets <n>'");

    // === Main Loop ===
    uint32_t tick = 0;
    while (1) {
        tick++;
        gpio_set_level(LED_PIN, (tick % 2) ? 1 : 0);
        if (tick % 100 == 0) {
            ESP_LOGI(TAG, "tick=%lu | pkts=%lu | LED=%d",
                     (unsigned long)tick, (unsigned long)packet_count, (int)(tick%2));
        }
        vTaskDelay(pdMS_TO_TICKS(1000));

        if (sniffer_active && difftime(time(NULL), last_csv_time) >= CSV_INTERVAL_SEC) {
            create_csv_file();
            last_csv_time = time(NULL);
        }
    }
}