/*
 * hijack_tx -- DMA hijack TX for ESP32-S3.
 *
 * Fires a blob TX on slot 0 with a dummy frame sized to match the real
 * frame, then swaps the DMA descriptor's buffer pointer before hardware
 * reads it. The blob's PLCP setup uses the correct length because the
 * dummy is the same size. No PLCP patching race.
 *
 * Sequence:
 *   1. Build dummy frame matching target length
 *   2. Fire blob TX (esp_wifi_80211_tx) on slot 0
 *   3. Spin-wait for slot 0 TXQ_CTRL bits [31:30] set
 *   4. Extract DMA descriptor from TXQ_CTRL[19:0]
 *   5. Swap desc->buf to our frame
 *   6. Wait for completion (bits [31:30] clear)
 *   7. Restore descriptor to prevent blob crash
 */

#include "hijack_tx.h"

#include <string.h>
#include <stdio.h>
#include "esp_heap_caps.h"
#include "esp_wifi.h"
#include "esp_mac.h"
#include "esp_timer.h"
#include "esp_rom_lldesc.h"
#include "xtensa/core-macros.h"

#define REG32(addr) (*(volatile uint32_t *)(addr))

#define S0_TXQ_CTRL   0x60033D08
#define S0_PMD        0x60034320

#define MAX_FRAME_SIZE 512
#define DMA_BASE       0x3FC00000

static uint8_t *s_dummy = NULL;
static uint8_t *s_frame_buf = NULL;
static uint8_t  s_own_mac[6];

static uint32_t s_ok_count = 0;
static uint32_t s_miss_count = 0;
static uint32_t s_fail_count = 0;

void hijack_init(void)
{
    if (s_frame_buf) return;

    s_frame_buf = heap_caps_calloc(1, MAX_FRAME_SIZE,
                                   MALLOC_CAP_DMA | MALLOC_CAP_INTERNAL);
    s_dummy = heap_caps_calloc(1, MAX_FRAME_SIZE,
                               MALLOC_CAP_DMA | MALLOC_CAP_INTERNAL);
    assert(s_frame_buf);
    assert(s_dummy);

    esp_read_mac(s_own_mac, ESP_MAC_WIFI_STA);
}

static void build_dummy(uint16_t len)
{
    memset(s_dummy, 0, len);
    s_dummy[0] = 0xD0;
    memcpy(s_dummy + 4, "\xFF\xFF\xFF\xFF\xFF\xFF", 6);
    memcpy(s_dummy + 10, s_own_mac, 6);
    memcpy(s_dummy + 16, s_own_mac, 6);
    s_dummy[24] = 127;
    s_dummy[25] = 0x52;
    s_dummy[26] = 0x4B;
    s_dummy[27] = 0x59;
}

int IRAM_ATTR hijack_send(const uint8_t *frame, uint16_t len)
{
    if (!s_frame_buf) return -1;
    if (len == 0 || len > MAX_FRAME_SIZE) return -1;

    memcpy(s_frame_buf, frame, len);

    build_dummy(len);

    esp_wifi_80211_tx(WIFI_IF_STA, s_dummy, len, false);

    uint32_t spin_start = xthal_get_ccount();
    uint32_t spin_limit = 5000 * 240;
    lldesc_t *desc = NULL;
    volatile const uint8_t *orig_buf = NULL;
    uint16_t orig_size = 0;
    uint16_t orig_len = 0;

    while ((xthal_get_ccount() - spin_start) < spin_limit) {
        uint32_t ctrl = REG32(S0_TXQ_CTRL);
        if (ctrl & 0xC0000000) {
            desc = (lldesc_t *)(DMA_BASE | (ctrl & 0x000FFFFF));

            orig_buf  = desc->buf;
            orig_size = desc->size;
            orig_len  = desc->length;

            desc->buf = (volatile const uint8_t *)s_frame_buf;
            __asm__ __volatile__("memw" ::: "memory");

            break;
        }
    }

    if (!desc) {
        s_miss_count++;
        return -1;
    }

    int result = -1;
    int64_t poll_end = esp_timer_get_time() + 200000;
    while (esp_timer_get_time() < poll_end) {
        if ((REG32(S0_TXQ_CTRL) & 0xC0000000) == 0) {
            uint32_t pmd = REG32(S0_PMD);
            uint8_t pmd_code = (pmd >> 12) & 0xF;
            result = (pmd_code == 0 || pmd_code == 4) ? 0 : (int)pmd_code;
            break;
        }
    }

    desc->buf    = orig_buf;
    desc->size   = orig_size;
    desc->length = orig_len;
    __asm__ __volatile__("memw" ::: "memory");

    if (result == 0)
        s_ok_count++;
    else
        s_fail_count++;

    if (((s_ok_count + s_fail_count) % 100) == 0) {
        uint32_t swap_us = (xthal_get_ccount() - spin_start) / 240;
        printf("hijack: ok=%lu fail=%lu miss=%lu swap=%luus pmd=%d\n",
               (unsigned long)s_ok_count,
               (unsigned long)s_fail_count,
               (unsigned long)s_miss_count,
               (unsigned long)swap_us,
               result);
    }

    return result;
}