#include "jlink.h"
#include "audio.h"
#include "config.h"
#include "stm32f7xx_hal.h"
#include <string.h>

/* ---- DMA circular RX buffer ---- */
#define JLINK_RX_BUF_LEN    128u   /* must be power-of-2 */
static uint8_t  s_rx_buf[JLINK_RX_BUF_LEN];
static uint32_t s_rx_tail = 0;     /* consumer index (byte already processed) */

/* ---- HAL handles ---- */
static UART_HandleTypeDef s_uart;
static DMA_HandleTypeDef  s_dma_rx;

/* ---- TX mutex ---- */
/*
 * Issue #522: USART1 IDLE interrupt (DMA RX) fires via HAL_UART_IRQHandler
 * mid-frame during polling HAL_UART_Transmit, resetting gState and causing
 * truncated/null-prefixed frames on the Jetson link.
 *
 * Fix: disable USART1_IRQn around every blocking TX so HAL_UART_IRQHandler
 * cannot modify gState while HAL_UART_Transmit is looping.  s_tx_busy guards
 * against any re-entrant caller (ESC debug, future paths).
 */
static volatile uint8_t s_tx_busy = 0;

static void jlink_tx_locked(uint8_t *buf, uint16_t len)
{
    if (s_tx_busy) return;          /* drop if already transmitting */
    s_tx_busy = 1u;
    HAL_NVIC_DisableIRQ(USART1_IRQn);
    HAL_UART_Transmit(&s_uart, buf, len, 5u);
    HAL_NVIC_EnableIRQ(USART1_IRQn);
    s_tx_busy = 0u;
}

/* ---- Volatile state ---- */
volatile JLinkState jlink_state;

/* ---- CRC16-XModem (poly 0x1021, init 0x0000) ---- */
static uint16_t crc16_xmodem(const uint8_t *data, uint16_t len)
{
    uint16_t crc = 0x0000u;
    for (uint16_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i] << 8;
        for (uint8_t b = 0; b < 8; b++) {
            if (crc & 0x8000u)
                crc = (crc << 1) ^ 0x1021u;
            else
                crc <<= 1;
        }
    }
    return crc;
}

/* ---- jlink_init() ---- */
void jlink_init(void)
{
    /* GPIO: PB6=TX AF7 (USART1_TX), PB7=RX AF7 (USART1_RX) */
    __HAL_RCC_GPIOB_CLK_ENABLE();
    GPIO_InitTypeDef gpio = {0};
    gpio.Pin       = GPIO_PIN_6 | GPIO_PIN_7;
    gpio.Mode      = GPIO_MODE_AF_PP;
    gpio.Pull      = GPIO_PULLUP;
    gpio.Speed     = GPIO_SPEED_FREQ_HIGH;
    gpio.Alternate = GPIO_AF7_USART1;
    HAL_GPIO_Init(GPIOB, &gpio);

    /* DMA2 Stream2 Channel4 — USART1_RX circular */
    __HAL_RCC_DMA2_CLK_ENABLE();
    s_dma_rx.Instance                 = DMA2_Stream2;
    s_dma_rx.Init.Channel             = DMA_CHANNEL_4;
    s_dma_rx.Init.Direction           = DMA_PERIPH_TO_MEMORY;
    s_dma_rx.Init.PeriphInc           = DMA_PINC_DISABLE;
    s_dma_rx.Init.MemInc              = DMA_MINC_ENABLE;
    s_dma_rx.Init.PeriphDataAlignment = DMA_PDATAALIGN_BYTE;
    s_dma_rx.Init.MemDataAlignment    = DMA_MDATAALIGN_BYTE;
    s_dma_rx.Init.Mode                = DMA_CIRCULAR;
    s_dma_rx.Init.Priority            = DMA_PRIORITY_MEDIUM;
    s_dma_rx.Init.FIFOMode            = DMA_FIFOMODE_DISABLE;
    HAL_DMA_Init(&s_dma_rx);
    __HAL_LINKDMA(&s_uart, hdmarx, s_dma_rx);

    /* USART1 at JLINK_BAUD (921600) */
    __HAL_RCC_USART1_CLK_ENABLE();
    s_uart.Instance          = USART1;
    s_uart.Init.BaudRate     = JLINK_BAUD;
    s_uart.Init.WordLength   = UART_WORDLENGTH_8B;
    s_uart.Init.StopBits     = UART_STOPBITS_1;
    s_uart.Init.Parity       = UART_PARITY_NONE;
    s_uart.Init.Mode         = UART_MODE_TX_RX;
    s_uart.Init.HwFlowCtl    = UART_HWCONTROL_NONE;
    s_uart.Init.OverSampling = UART_OVERSAMPLING_16;
    HAL_UART_Init(&s_uart);

    /* Enable USART1 IDLE interrupt for circular buffer draining */
    __HAL_UART_ENABLE_IT(&s_uart, UART_IT_IDLE);
    HAL_NVIC_SetPriority(USART1_IRQn, 6, 0);
    HAL_NVIC_EnableIRQ(USART1_IRQn);

    /* DMA2_Stream2 IRQ (for error handling) */
    HAL_NVIC_SetPriority(DMA2_Stream2_IRQn, 7, 0);
    HAL_NVIC_EnableIRQ(DMA2_Stream2_IRQn);

    /* Start circular DMA RX — never stops */
    HAL_UART_Receive_DMA(&s_uart, s_rx_buf, JLINK_RX_BUF_LEN);

    memset((void *)&jlink_state, 0, sizeof(jlink_state));
    s_rx_tail = 0;
}

/* ---- IRQ handlers ---- */
void USART1_IRQHandler(void)
{
    /* Clear IDLE flag by reading SR then DR */
    if (__HAL_UART_GET_FLAG(&s_uart, UART_FLAG_IDLE)) {
        __HAL_UART_CLEAR_IDLEFLAG(&s_uart);
        /* jlink_process() drains the buffer from main loop — no work here */
    }
    HAL_UART_IRQHandler(&s_uart);
}

void DMA2_Stream2_IRQHandler(void)
{
    HAL_DMA_IRQHandler(&s_dma_rx);
}

/* ---- jlink_is_active() ---- */
bool jlink_is_active(uint32_t now_ms)
{
    if (jlink_state.last_rx_ms == 0u) return false;
    return (now_ms - jlink_state.last_rx_ms) < JLINK_HB_TIMEOUT_MS;
}

/* ---- Frame dispatch ---- */
static void dispatch(const uint8_t *payload, uint8_t cmd, uint8_t plen)
{
    /* Update heartbeat timestamp on every valid frame */
    jlink_state.last_rx_ms = HAL_GetTick();

    switch (cmd) {

    case JLINK_CMD_HEARTBEAT:
        /* Heartbeat only — no payload action needed */
        break;

    case JLINK_CMD_DRIVE:
        if (plen == 4u) {
            int16_t spd, str;
            memcpy(&spd, payload,     2);
            memcpy(&str, payload + 2, 2);
            /* Clamp to ±1000 */
            if (spd >  1000) spd =  1000;
            if (spd < -1000) spd = -1000;
            if (str >  1000) str =  1000;
            if (str < -1000) str = -1000;
            jlink_state.speed = spd;
            jlink_state.steer = str;
        }
        break;

    case JLINK_CMD_ARM:
        jlink_state.arm_req = 1u;
        break;

    case JLINK_CMD_DISARM:
        jlink_state.disarm_req = 1u;
        break;

    case JLINK_CMD_PID_SET:
        if (plen == 12u) {
            float kp, ki, kd;
            memcpy(&kp, payload,     4);
            memcpy(&ki, payload + 4, 4);
            memcpy(&kd, payload + 8, 4);
            /* Sanity bounds — same as USB CDC PID handler in main.c */
            if (kp >= 0.0f && kp <= 500.0f) jlink_state.pid_kp = kp;
            if (ki >= 0.0f && ki <= 50.0f)  jlink_state.pid_ki = ki;
            if (kd >= 0.0f && kd <= 50.0f)  jlink_state.pid_kd = kd;
            jlink_state.pid_updated = 1u;
        }
        break;

    case JLINK_CMD_DFU_ENTER:
        /* Payload-less; main loop checks armed state before calling ota_enter_dfu() */
        jlink_state.dfu_req = 1u;
        break;

    case JLINK_CMD_ESTOP:
        jlink_state.estop_req = 1u;
        break;

    case JLINK_CMD_AUDIO:
        /* Payload: int16 PCM samples, little-endian, 1..126 samples (2..252 bytes) */
        if (plen >= 2u && (plen & 1u) == 0u) {
            audio_write_pcm((const int16_t *)payload, plen / 2u);
        }
        break;

    case JLINK_CMD_SLEEP:
        /* Payload-less; main loop calls power_mgmt_request_sleep() */
        jlink_state.sleep_req = 1u;
        break;

    case JLINK_CMD_PID_SAVE:
        /* Payload-less; main loop calls pid_flash_save() (Issue #531) */
        jlink_state.pid_save_req = 1u;
        break;

    default:
        break;
    }
}

/* ---- jlink_process() — call from main loop every tick ---- */
/*
 * Parser state machine.
 * Frame: [STX][LEN][CMD][PAYLOAD 0..LEN-1][CRC_hi][CRC_lo][ETX]
 * LEN = count of CMD + PAYLOAD bytes (1..253).
 * CRC16-XModem over CMD+PAYLOAD.
 * Maximum payload = 253 - 1 = 252 bytes (LEN field is 1 byte, max 0xFF=255,
 * but we cap at 64 for safety).
 */
#define JLINK_MAX_PAYLOAD   252u  /* enlarged for AUDIO chunks (126 × int16) */

typedef enum {
    PS_WAIT_STX = 0,
    PS_WAIT_LEN,
    PS_WAIT_DATA,   /* receiving CMD + PAYLOAD (len bytes total) */
    PS_WAIT_CRC_HI,
    PS_WAIT_CRC_LO,
    PS_WAIT_ETX,
} ParseState;

void jlink_process(void)
{
    static ParseState s_state   = PS_WAIT_STX;
    static uint8_t    s_len     = 0;    /* expected CMD+PAYLOAD length */
    static uint8_t    s_count   = 0;    /* bytes received so far in PS_WAIT_DATA */
    static uint8_t    s_frame[JLINK_MAX_PAYLOAD + 1u]; /* [0]=CMD, [1..]=PAYLOAD */
    static uint8_t    s_crc_hi  = 0;

    /* Compute how many bytes the DMA has written since last drain */
    uint32_t head = JLINK_RX_BUF_LEN - __HAL_DMA_GET_COUNTER(&s_dma_rx);
    uint32_t bytes = (head - s_rx_tail) & (JLINK_RX_BUF_LEN - 1u);

    for (uint32_t i = 0; i < bytes; i++) {
        uint8_t b = s_rx_buf[s_rx_tail];
        s_rx_tail = (s_rx_tail + 1u) & (JLINK_RX_BUF_LEN - 1u);

        switch (s_state) {
        case PS_WAIT_STX:
            if (b == JLINK_STX) s_state = PS_WAIT_LEN;
            break;

        case PS_WAIT_LEN:
            if (b == 0u || b > JLINK_MAX_PAYLOAD + 1u) {
                /* Invalid length — resync */
                s_state = PS_WAIT_STX;
            } else {
                s_len   = b;
                s_count = 0;
                s_state = PS_WAIT_DATA;
            }
            break;

        case PS_WAIT_DATA:
            s_frame[s_count++] = b;
            if (s_count == s_len) s_state = PS_WAIT_CRC_HI;
            break;

        case PS_WAIT_CRC_HI:
            s_crc_hi = b;
            s_state  = PS_WAIT_CRC_LO;
            break;

        case PS_WAIT_CRC_LO: {
            uint16_t rx_crc  = ((uint16_t)s_crc_hi << 8) | b;
            uint16_t calc_crc = crc16_xmodem(s_frame, s_len);
            if (rx_crc == calc_crc)
                s_state = PS_WAIT_ETX;
            else
                s_state = PS_WAIT_STX;   /* CRC mismatch — drop */
            break;
        }

        case PS_WAIT_ETX:
            if (b == JLINK_ETX) {
                /* Valid frame: s_frame[0]=CMD, s_frame[1..s_len-1]=PAYLOAD */
                dispatch(s_frame + 1, s_frame[0], s_len - 1u);
            }
            /* Either way, go back to idle (resync on bad ETX) */
            s_state = PS_WAIT_STX;
            break;
        }
    }
}

/* ---- jlink_send_telemetry() ---- */
void jlink_send_telemetry(const jlink_tlm_status_t *status)
{
    /*
     * Frame: [STX][LEN][0x80][20 bytes STATUS][CRC_hi][CRC_lo][ETX]
     * LEN = 1 (CMD) + 20 (payload) = 21
     * Total frame length = 1+1+1+20+2+1 = 26 bytes
     * At 921600 baud (10 bits/byte): 26×10/921600 ≈ 0.28ms — safe to block.
     */
    static uint8_t frame[26];
    const uint8_t  plen = (uint8_t)sizeof(jlink_tlm_status_t);  /* 20 */
    const uint8_t  len  = 1u + plen;                            /* 21 */

    frame[0] = JLINK_STX;
    frame[1] = len;
    frame[2] = JLINK_TLM_STATUS;
    memcpy(&frame[3], status, plen);

    uint16_t crc = crc16_xmodem(&frame[2], len);  /* over CMD + PAYLOAD */
    frame[3 + plen]     = (uint8_t)(crc >> 8);
    frame[3 + plen + 1] = (uint8_t)(crc & 0xFFu);
    frame[3 + plen + 2] = JLINK_ETX;

    jlink_tx_locked(frame, sizeof(frame));
}

/* ---- jlink_send_power_telemetry() ---- */
void jlink_send_power_telemetry(const jlink_tlm_power_t *power)
{
    /*
     * Frame: [STX][LEN][0x81][11 bytes POWER][CRC_hi][CRC_lo][ETX]
     * LEN = 1 (CMD) + 11 (payload) = 12; total = 17 bytes
     * At 921600 baud: 17×10/921600 ≈ 0.18 ms — safe to block.
     */
    static uint8_t frame[17];
    const uint8_t  plen = (uint8_t)sizeof(jlink_tlm_power_t);  /* 11 */
    const uint8_t  len  = 1u + plen;                            /* 12 */

    frame[0] = JLINK_STX;
    frame[1] = len;
    frame[2] = JLINK_TLM_POWER;
    memcpy(&frame[3], power, plen);

    uint16_t crc = crc16_xmodem(&frame[2], len);
    frame[3 + plen]     = (uint8_t)(crc >> 8);
    frame[3 + plen + 1] = (uint8_t)(crc & 0xFFu);
    frame[3 + plen + 2] = JLINK_ETX;

    jlink_tx_locked(frame, sizeof(frame));
}

/* ---- jlink_send_pid_result() -- Issue #531 ---- */
void jlink_send_pid_result(const jlink_tlm_pid_result_t *result)
{
    /*
     * Frame: [STX][LEN][0x83][13 bytes PID_RESULT][CRC_hi][CRC_lo][ETX]
     * LEN = 1 (CMD) + 13 (payload) = 14; total frame = 19 bytes.
     * At 921600 baud: 19x10/921600 ~0.21 ms -- safe to block.
     */
    static uint8_t frame_pid[19];
    const uint8_t  plen = (uint8_t)sizeof(jlink_tlm_pid_result_t);  /* 13 */
    const uint8_t  flen = 1u + plen;                                 /* 14 */

    frame_pid[0] = JLINK_STX;
    frame_pid[1] = flen;
    frame_pid[2] = JLINK_TLM_PID_RESULT;
    memcpy(&frame_pid[3], result, plen);

    uint16_t crc_pid = crc16_xmodem(&frame_pid[2], flen);
    frame_pid[3 + plen]     = (uint8_t)(crc_pid >> 8);
    frame_pid[3 + plen + 1] = (uint8_t)(crc_pid & 0xFFu);
    frame_pid[3 + plen + 2] = JLINK_ETX;

    jlink_tx_locked(frame_pid, sizeof(frame_pid));
}

/* ---- jlink_send_battery_telemetry() — Issue #533 ---- */
void jlink_send_battery_telemetry(const jlink_tlm_battery_t *batt)
{
    /*
     * Frame: [STX][LEN][0x82][10 bytes BATTERY][CRC_hi][CRC_lo][ETX]
     * LEN = 1 (CMD) + 10 (payload) = 11; total = 16 bytes
     * At 921600 baud: 16×10/921600 ≈ 0.17 ms — safe to block.
     */
    static uint8_t frame[16];
    const uint8_t  plen = (uint8_t)sizeof(jlink_tlm_battery_t);  /* 10 */
    const uint8_t  len  = 1u + plen;                              /* 11 */

    frame[0] = JLINK_STX;
    frame[1] = len;
    frame[2] = JLINK_TLM_BATTERY;
    memcpy(&frame[3], batt, plen);

    uint16_t crc = crc16_xmodem(&frame[2], len);
    frame[3 + plen]     = (uint8_t)(crc >> 8);
    frame[3 + plen + 1] = (uint8_t)(crc & 0xFFu);
    frame[3 + plen + 2] = JLINK_ETX;

    jlink_tx_locked(frame, sizeof(frame));
}