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

/* ================================================================
 * Buffer layout
 * ================================================================
 * AUDIO_BUF_HALF samples per DMA half (config.h: 441 at 22050 Hz → 20 ms).
 * DMA runs in circular mode over 2 halves; TxHalfCplt refills [0] and
 * TxCplt refills [AUDIO_BUF_HALF].  Both callbacks simply call fill_half().
 */
#define AUDIO_BUF_SIZE    (AUDIO_BUF_HALF * 2u)

/* DMA buffer: must be in non-cached SRAM or flushed before DMA access.
 * Placed in SRAM1 (default .bss section on STM32F722 — below 512 KB).
 * The I-cache / D-cache is on by default; DMA1 is an AHB master that
 * bypasses cache, so we keep this buffer in DTCM-uncacheable SRAM. */
static int16_t s_dma_buf[AUDIO_BUF_SIZE];

/* ================================================================
 * PCM FIFO — Jetson audio (main-loop writer, ISR reader)
 * ================================================================
 * Single-producer / single-consumer lock-free ring buffer.
 * Power-of-2 size so wrap uses bitwise AND (safe across ISR boundary).
 * 4096 samples ≈ 185 ms at 22050 Hz — enough to absorb JLink jitter.
 */
#define PCM_FIFO_SIZE   4096u
#define PCM_FIFO_MASK   (PCM_FIFO_SIZE - 1u)

static int16_t           s_pcm_fifo[PCM_FIFO_SIZE];
static volatile uint16_t s_pcm_rd = 0;   /* consumer (ISR advances) */
static volatile uint16_t s_pcm_wr = 0;   /* producer (main loop advances) */

/* ================================================================
 * Tone sequencer
 * ================================================================
 * Each AudioTone maps to a const ToneStep array.  Steps are played
 * sequentially; a gap_ms of 0 means no silence between steps.
 * audio_tick() in the main loop advances the state machine and
 * writes the volatile active-tone params read by the ISR fill path.
 */
typedef struct {
    uint16_t freq_hz;
    uint16_t dur_ms;
    uint16_t gap_ms;   /* silence after this step */
} ToneStep;

/* ---- Tone definitions ---- */
static const ToneStep s_def_beep_short[]  = {{880,  100, 0}};
static const ToneStep s_def_beep_long[]   = {{880,  500, 0}};
static const ToneStep s_def_startup[]     = {{523,  120, 60},
                                              {659,  120, 60},
                                              {784,  200, 0}};
static const ToneStep s_def_arm[]         = {{880,   80, 60},
                                              {1047, 100, 0}};
static const ToneStep s_def_disarm[]      = {{880,   80, 60},
                                              {659,  100, 0}};
static const ToneStep s_def_fault[]       = {{200,  500, 0}};

typedef struct {
    const ToneStep *steps;
    uint8_t n_steps;
} ToneDef;

static const ToneDef s_tone_defs[AUDIO_TONE_COUNT] = {
    [AUDIO_TONE_BEEP_SHORT] = {s_def_beep_short, 1},
    [AUDIO_TONE_BEEP_LONG]  = {s_def_beep_long,  1},
    [AUDIO_TONE_STARTUP]    = {s_def_startup,     3},
    [AUDIO_TONE_ARM]        = {s_def_arm,         2},
    [AUDIO_TONE_DISARM]     = {s_def_disarm,      2},
    [AUDIO_TONE_FAULT]      = {s_def_fault,       1},
};

/* Active tone queue */
#define TONE_QUEUE_DEPTH  4u

typedef struct {
    const ToneDef *def;
    uint8_t  step;           /* current step index within def */
    bool     in_gap;         /* true while playing inter-step silence */
    uint32_t step_end_ms;    /* abs HAL_GetTick() when step/gap expires */
} ToneSeq;

static ToneSeq   s_tone_q[TONE_QUEUE_DEPTH];
static uint8_t   s_tq_head = 0;   /* consumer (audio_tick reads) */
static uint8_t   s_tq_tail = 0;   /* producer (audio_play_tone writes) */

/* Volatile parameters written by audio_tick(), read by ISR fill_half() */
static volatile uint16_t s_active_freq  = 0;   /* 0 = silence / gap */
static volatile uint32_t s_active_phase = 0;   /* sample phase counter */

/* ================================================================
 * Volume
 * ================================================================ */
static uint8_t s_volume = AUDIO_VOLUME_DEFAULT;  /* 0–100 */

/* ================================================================
 * HAL handles
 * ================================================================ */
static I2S_HandleTypeDef s_i2s;
static DMA_HandleTypeDef s_dma_tx;

/* ================================================================
 * fill_half() — called from ISR, O(AUDIO_BUF_HALF)
 * ================================================================
 * Priority: PCM FIFO (Jetson TTS) > square-wave tone > silence.
 * Volume applied via integer scaling — no float in ISR.
 */
static void fill_half(int16_t *buf, uint16_t n)
{
    uint16_t rd    = s_pcm_rd;
    uint16_t wr    = s_pcm_wr;
    uint16_t avail = (uint16_t)((wr - rd) & PCM_FIFO_MASK);

    if (avail >= n) {
        /* ---- Drain Jetson PCM FIFO ---- */
        for (uint16_t i = 0; i < n; i++) {
            int32_t s = (int32_t)s_pcm_fifo[rd] * (int32_t)s_volume / 100;
            buf[i] = (s > 32767) ? (int16_t)32767 :
                     (s < -32768) ? (int16_t)-32768 : (int16_t)s;
            rd = (rd + 1u) & PCM_FIFO_MASK;
        }
        s_pcm_rd = rd;

    } else if (s_active_freq) {
        /* ---- Square wave tone generator ---- */
        uint32_t half_p = (uint32_t)AUDIO_SAMPLE_RATE / (2u * s_active_freq);
        int16_t  amp    = (int16_t)(16384u * (uint32_t)s_volume / 100u);
        uint32_t ph     = s_active_phase;
        uint32_t period = 2u * half_p;
        for (uint16_t i = 0; i < n; i++) {
            buf[i] = ((ph % period) < half_p) ? amp : (int16_t)(-amp);
            ph++;
        }
        s_active_phase = ph;

    } else {
        /* ---- Silence ---- */
        memset(buf, 0, (size_t)(n * 2u));
    }
}

/* ================================================================
 * DMA callbacks (ISR context)
 * ================================================================ */
void HAL_I2S_TxHalfCpltCallback(I2S_HandleTypeDef *hi2s)
{
    if (hi2s->Instance == SPI3)
        fill_half(s_dma_buf, AUDIO_BUF_HALF);
}

void HAL_I2S_TxCpltCallback(I2S_HandleTypeDef *hi2s)
{
    if (hi2s->Instance == SPI3)
        fill_half(s_dma_buf + AUDIO_BUF_HALF, AUDIO_BUF_HALF);
}

/* DMA1 Stream7 IRQ (SPI3/I2S3 TX) */
void DMA1_Stream7_IRQHandler(void)
{
    HAL_DMA_IRQHandler(&s_dma_tx);
}

/* ================================================================
 * audio_init()
 * ================================================================ */
void audio_init(void)
{
    /* ---- PLLI2S: N=192, R=2 → 96 MHz I2S clock ----
     * With SPI3/I2S3 and I2S_DATAFORMAT_16B (16-bit, 32-bit frame slot):
     *   FS = 96 MHz / (32 × 2 × I2SDIV) where HAL picks I2SDIV = 68
     *   → 96 000 000 / (32 × 2 × 68) = 22 058 Hz  (< 0.04 % error)
     */
    RCC_PeriphCLKInitTypeDef pclk = {0};
    pclk.PeriphClockSelection = RCC_PERIPHCLK_I2S;
    pclk.I2sClockSelection    = RCC_I2SCLKSOURCE_PLLI2S;
    pclk.PLLI2S.PLLI2SN       = 192;   /* VCO = (HSE/PLLM) × 192 = 192 MHz */
    pclk.PLLI2S.PLLI2SR       = 2;     /* I2S clock = 96 MHz */
    HAL_RCCEx_PeriphCLKConfig(&pclk);

    /* ---- GPIO ---- */
    __HAL_RCC_GPIOA_CLK_ENABLE();
    __HAL_RCC_GPIOB_CLK_ENABLE();
    __HAL_RCC_GPIOC_CLK_ENABLE();

    GPIO_InitTypeDef gpio = {0};

    /* PA15: I2S3_WS (LRCLK), AF6 */
    gpio.Pin       = AUDIO_LRCK_PIN;
    gpio.Mode      = GPIO_MODE_AF_PP;
    gpio.Pull      = GPIO_NOPULL;
    gpio.Speed     = GPIO_SPEED_FREQ_VERY_HIGH;
    gpio.Alternate = GPIO_AF6_SPI3;
    HAL_GPIO_Init(AUDIO_LRCK_PORT, &gpio);

    /* PC10: I2S3_CK (BCLK), AF6 */
    gpio.Pin       = AUDIO_BCLK_PIN;
    HAL_GPIO_Init(AUDIO_BCLK_PORT, &gpio);

    /* PB5: I2S3_SD (DIN), AF6 */
    gpio.Pin       = AUDIO_DOUT_PIN;
    HAL_GPIO_Init(AUDIO_DOUT_PORT, &gpio);

    /* PC5: AUDIO_MUTE GPIO output — drive low (muted) initially */
    gpio.Pin       = AUDIO_MUTE_PIN;
    gpio.Mode      = GPIO_MODE_OUTPUT_PP;
    gpio.Pull      = GPIO_NOPULL;
    gpio.Speed     = GPIO_SPEED_FREQ_LOW;
    gpio.Alternate = 0;
    HAL_GPIO_Init(AUDIO_MUTE_PORT, &gpio);
    HAL_GPIO_WritePin(AUDIO_MUTE_PORT, AUDIO_MUTE_PIN, GPIO_PIN_RESET);

    /* ---- DMA1 Stream7 Channel0 (SPI3/I2S3 TX) ---- */
    __HAL_RCC_DMA1_CLK_ENABLE();
    s_dma_tx.Instance                 = DMA1_Stream7;
    s_dma_tx.Init.Channel             = DMA_CHANNEL_0;
    s_dma_tx.Init.Direction           = DMA_MEMORY_TO_PERIPH;
    s_dma_tx.Init.PeriphInc           = DMA_PINC_DISABLE;
    s_dma_tx.Init.MemInc              = DMA_MINC_ENABLE;
    s_dma_tx.Init.PeriphDataAlignment = DMA_PDATAALIGN_HALFWORD;
    s_dma_tx.Init.MemDataAlignment    = DMA_MDATAALIGN_HALFWORD;
    s_dma_tx.Init.Mode                = DMA_CIRCULAR;
    s_dma_tx.Init.Priority            = DMA_PRIORITY_MEDIUM;
    s_dma_tx.Init.FIFOMode            = DMA_FIFOMODE_DISABLE;
    HAL_DMA_Init(&s_dma_tx);
    __HAL_LINKDMA(&s_i2s, hdmatx, s_dma_tx);

    HAL_NVIC_SetPriority(DMA1_Stream7_IRQn, 7, 0);
    HAL_NVIC_EnableIRQ(DMA1_Stream7_IRQn);

    /* ---- SPI3 in I2S3 master TX mode ---- */
    __HAL_RCC_SPI3_CLK_ENABLE();
    s_i2s.Instance         = SPI3;
    s_i2s.Init.Mode        = I2S_MODE_MASTER_TX;
    s_i2s.Init.Standard    = I2S_STANDARD_PHILIPS;
    s_i2s.Init.DataFormat  = I2S_DATAFORMAT_16B;
    s_i2s.Init.MCLKOutput  = I2S_MCLKOUTPUT_DISABLE;
    s_i2s.Init.AudioFreq   = I2S_AUDIOFREQ_22K;
    s_i2s.Init.CPOL        = I2S_CPOL_LOW;
    s_i2s.Init.ClockSource = I2S_CLOCK_PLL;
    HAL_I2S_Init(&s_i2s);

    /* Pre-fill with silence and start circular DMA TX */
    memset(s_dma_buf, 0, sizeof(s_dma_buf));
    HAL_I2S_Transmit_DMA(&s_i2s, (uint16_t *)s_dma_buf, AUDIO_BUF_SIZE);

    /* Unmute amp after DMA is running — avoids start-up click */
    HAL_GPIO_WritePin(AUDIO_MUTE_PORT, AUDIO_MUTE_PIN, GPIO_PIN_SET);
}

/* ================================================================
 * Public API
 * ================================================================ */

void audio_mute(bool active)
{
    HAL_GPIO_WritePin(AUDIO_MUTE_PORT, AUDIO_MUTE_PIN,
                      active ? GPIO_PIN_SET : GPIO_PIN_RESET);
}

void audio_set_volume(uint8_t vol)
{
    s_volume = (vol > 100u) ? 100u : vol;
}

bool audio_play_tone(AudioTone tone)
{
    if (tone >= AUDIO_TONE_COUNT) return false;

    uint8_t next = (s_tq_tail + 1u) % TONE_QUEUE_DEPTH;
    if (next == s_tq_head) return false;  /* queue full */

    s_tone_q[s_tq_tail].def         = &s_tone_defs[tone];
    s_tone_q[s_tq_tail].step        = 0;
    s_tone_q[s_tq_tail].in_gap      = false;
    s_tone_q[s_tq_tail].step_end_ms = 0;  /* audio_tick() sets this on first run */
    s_tq_tail = next;
    return true;
}

uint16_t audio_write_pcm(const int16_t *samples, uint16_t n)
{
    uint16_t wr    = s_pcm_wr;
    uint16_t rd    = s_pcm_rd;
    uint16_t space = (uint16_t)((rd - wr - 1u) & PCM_FIFO_MASK);
    uint16_t accept = (n < space) ? n : space;

    for (uint16_t i = 0; i < accept; i++) {
        s_pcm_fifo[wr] = samples[i];
        wr = (wr + 1u) & PCM_FIFO_MASK;
    }
    s_pcm_wr = wr;
    return accept;
}

void audio_tick(uint32_t now_ms)
{
    /* Nothing to do if queue is empty */
    if (s_tq_head == s_tq_tail) {
        s_active_freq = 0;
        return;
    }

    ToneSeq *seq = &s_tone_q[s_tq_head];

    /* First call for this sequence entry: arm the first step */
    if (seq->step_end_ms == 0u) {
        const ToneStep *st = &seq->def->steps[0];
        seq->in_gap      = false;
        seq->step_end_ms = now_ms + st->dur_ms;
        s_active_freq    = st->freq_hz;
        s_active_phase   = 0;
        return;
    }

    /* Step / gap still running */
    if (now_ms < seq->step_end_ms) return;

    /* Current step or gap has expired */
    const ToneStep *st = &seq->def->steps[seq->step];

    if (!seq->in_gap && st->gap_ms) {
        /* Transition: tone → inter-step gap (silence) */
        seq->in_gap      = true;
        seq->step_end_ms = now_ms + st->gap_ms;
        s_active_freq    = 0;
        return;
    }

    /* Advance to next step */
    seq->step++;
    seq->in_gap = false;

    if (seq->step >= seq->def->n_steps) {
        /* Sequence complete — pop from queue */
        s_tq_head = (s_tq_head + 1u) % TONE_QUEUE_DEPTH;
        s_active_freq = 0;
        return;
    }

    /* Start next step */
    st               = &seq->def->steps[seq->step];
    seq->step_end_ms = now_ms + st->dur_ms;
    s_active_freq    = st->freq_hz;
    s_active_phase   = 0;
}

bool audio_is_playing(void)
{
    return (s_i2s.State == HAL_I2S_STATE_BUSY_TX);
}
