/*
 * DMA hijack TX: 440Hz sine wave via slot 0 buffer swap.
 * 48kHz 16-bit stereo interleaved. 48 samples/pkt = 192 bytes audio.
 * 1000 packets/sec (1ms period). Uses ricky's original hijack_send.
 */

#include <string.h>
#include <math.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_timer.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "nvs_flash.h"
#include "esp_netif.h"
#include "esp_event.h"

#include "hijack_tx.h"

#define TAG "hijack-tx"

static const uint8_t BCAST[] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
static const uint8_t OUI_RKY[] = {0x52, 0x4B, 0x59};

#define CHANNEL           1
#define SAMPLE_RATE       48000
#define SINE_FREQ         440
#define SAMPLES_PER_PKT   48
#define BYTES_PER_SAMPLE  4
#define AUDIO_PAYLOAD     (SAMPLES_PER_PKT * BYTES_PER_SAMPLE)

#define TX_PERIOD_US      1000
#define STATS_INTERVAL_MS 5000

typedef struct __attribute__((packed)) {
    uint32_t frame_seq;
    uint8_t  sub_index;
    uint8_t  pkt_type;
    uint16_t payload_len;
    uint16_t burst_width;
} rky_hdr_t;

#define MAC_HDR_LEN    24
#define ACTION_HDR_LEN 4
#define RKY_HDR_LEN    sizeof(rky_hdr_t)
#define FRAME_OVERHEAD (MAC_HDR_LEN + ACTION_HDR_LEN + RKY_HDR_LEN)
#define MAX_FRAME_LEN  (FRAME_OVERHEAD + AUDIO_PAYLOAD)

/* pre-computed sine LUT: 4800 samples = 100 full 440Hz cycles at 48kHz */
#define SINE_TABLE_LEN 4800
static int16_t sine_table[SINE_TABLE_LEN];
static uint32_t s_sample_offset = 0;

static void init_sine_table(void)
{
    for (int i = 0; i < SINE_TABLE_LEN; i++) {
        float t = (float)i / (float)SAMPLE_RATE;
        sine_table[i] = (int16_t)(sinf(2.0f * M_PI * SINE_FREQ * t) * 4096.0f);
    }
}

static void gen_sine(int16_t *buf, int n_samples)
{
    for (int i = 0; i < n_samples; i++) {
        int16_t v = sine_table[(s_sample_offset + i) % SINE_TABLE_LEN];
        buf[i * 2]     = v;
        buf[i * 2 + 1] = v;
    }
    s_sample_offset = (s_sample_offset + n_samples) % SINE_TABLE_LEN;
}

static int build_frame(uint8_t *buf, const uint8_t *src,
                       uint32_t seq, const uint8_t *payload, uint16_t plen)
{
    memset(buf, 0, MAC_HDR_LEN);
    buf[0] = 0xD0;
    memcpy(buf + 4, BCAST, 6);
    memcpy(buf + 10, src, 6);
    memcpy(buf + 16, src, 6);

    int off = MAC_HDR_LEN;
    buf[off++] = 127;
    buf[off++] = OUI_RKY[0];
    buf[off++] = OUI_RKY[1];
    buf[off++] = OUI_RKY[2];

    rky_hdr_t *hdr = (rky_hdr_t *)(buf + off);
    hdr->frame_seq   = seq;
    hdr->sub_index   = 0;
    hdr->pkt_type    = 0;
    hdr->payload_len = plen;
    hdr->burst_width = 1;
    off += RKY_HDR_LEN;

    if (payload && plen > 0) {
        memcpy(buf + off, payload, plen);
        off += plen;
    }
    return off;
}

static void tx_task(void *arg)
{
    uint8_t mac[6];
    esp_read_mac(mac, ESP_MAC_WIFI_STA);

    init_sine_table();
    hijack_init();

    /* prime blob TX engine -- required before hijack_send works */
    {
        uint8_t dummy[64] = {0};
        dummy[0] = 0xD0;
        memcpy(dummy + 4, BCAST, 6);
        memcpy(dummy + 10, mac, 6);
        memcpy(dummy + 16, mac, 6);
        dummy[24] = 127;
        dummy[25] = OUI_RKY[0];
        dummy[26] = OUI_RKY[1];
        dummy[27] = OUI_RKY[2];
        esp_wifi_80211_tx(WIFI_IF_STA, dummy, sizeof(dummy), false);
        vTaskDelay(pdMS_TO_TICKS(10));
    }

    ESP_LOGI(TAG, "TX start: %dHz sine, %dkHz/16-bit stereo, %d bytes/pkt, period=%dus",
             SINE_FREQ, SAMPLE_RATE / 1000, AUDIO_PAYLOAD, TX_PERIOD_US);

    uint8_t frame[MAX_FRAME_LEN];
    int16_t audio_buf[SAMPLES_PER_PKT * 2];

    uint32_t seq = 0;
    uint32_t ok = 0, err = 0, miss = 0;
    int64_t dur_sum = 0, dur_min = INT64_MAX, dur_max = 0;
    uint32_t dur_n = 0;
    int64_t last_stats = esp_timer_get_time();
    int64_t next_tx = esp_timer_get_time();

    while (1) {
        while (esp_timer_get_time() < next_tx) { }

        next_tx += TX_PERIOD_US;
        int64_t now = esp_timer_get_time();
        if (now >= next_tx)
            next_tx = now + TX_PERIOD_US;

        int64_t t0 = esp_timer_get_time();

        gen_sine(audio_buf, SAMPLES_PER_PKT);
        int flen = build_frame(frame, mac, seq,
                               (uint8_t *)audio_buf, AUDIO_PAYLOAD);

        esp_err_t ret = esp_wifi_80211_tx(WIFI_IF_STA, frame, flen, false);
        if (ret == ESP_OK)  ok++;
        else                err++;

        seq++;

        int64_t dur = esp_timer_get_time() - t0;
        dur_sum += dur;
        dur_n++;
        if (dur < dur_min) dur_min = dur;
        if (dur > dur_max) dur_max = dur;

        if ((esp_timer_get_time() - last_stats) >= (int64_t)STATS_INTERVAL_MS * 1000) {
            double avg = dur_n > 0 ? (double)dur_sum / dur_n : 0;
            double rate = dur_n > 0 ? (double)dur_n / ((double)STATS_INTERVAL_MS / 1000.0) : 0;
            ESP_LOGI(TAG, "seq=%lu ok=%lu err=%lu miss=%lu "
                     "dur avg=%.0f min=%lld max=%lld pps=%.0f",
                     (unsigned long)seq,
                     (unsigned long)ok, (unsigned long)err,
                     (unsigned long)miss,
                     avg, (long long)dur_min, (long long)dur_max, rate);

            ok = err = miss = dur_n = 0;
            dur_sum = 0;
            dur_min = INT64_MAX;
            dur_max = 0;
            last_stats = esp_timer_get_time();
        }
    }
}

static void wifi_init(void)
{
    ESP_ERROR_CHECK(nvs_flash_init());
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));

    wifi_config_t sta_cfg = {0};
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &sta_cfg));
    ESP_ERROR_CHECK(esp_wifi_set_protocol(WIFI_IF_STA,
        WIFI_PROTOCOL_11B | WIFI_PROTOCOL_11G | WIFI_PROTOCOL_11N));

    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_ERROR_CHECK(esp_wifi_set_channel(CHANNEL, WIFI_SECOND_CHAN_NONE));
    ESP_ERROR_CHECK(esp_wifi_config_80211_tx_rate(WIFI_IF_STA, WIFI_PHY_RATE_54M));
    ESP_ERROR_CHECK(esp_wifi_set_ps(WIFI_PS_NONE));
}

void app_main(void)
{
    wifi_init();

    ESP_LOGI(TAG, "440Hz sine TX via DMA hijack (ricky original)");
    xTaskCreatePinnedToCore(tx_task, "tx", 8192, NULL,
                            configMAX_PRIORITIES - 1, NULL, 0);
}