/* Audio: gong synthesis, mic capture -> gateway, gateway PCM -> speaker. */

#include <math.h>
#include <string.h>

#include "esp_heap_caps.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/task.h"
#include "bsp/esp-bsp.h"
#include "esp_codec_dev.h"

#include "esp_afe_config.h"
#include "esp_afe_sr_iface.h"
#include "esp_afe_sr_models.h"
#include "esp_wn_models.h"

#include "desklock.h"

static const char *TAG = "audio";

/* end an utterance after this much post-speech silence (AFE VAD) */
#define VAD_END_SILENCE_MS 800

#define RATE 16000
#define CAPTURE_CHUNK 3200 /* 100 ms */
#define REPLY_MAX (RATE * 2 * 60) /* 60 s of reply audio */

#define GONG_SECONDS 5
#define GONG_SAMPLES (RATE * GONG_SECONDS)
#define GONG_VOLUME 75
#define GONG_PEAK 14000.0f

static esp_codec_dev_handle_t s_spk;
static esp_codec_dev_handle_t s_mic;
static int16_t *s_gong;

static uint8_t *s_reply;
static volatile size_t s_reply_len;
static volatile bool s_playing;

/* wake word + utterance (esp-sr AFE) */
static const esp_afe_sr_iface_t *s_afe;
static esp_afe_sr_data_t *s_afe_data;
static volatile bool s_streaming;   /* detect_task streams AFE audio upstream while true */
static int16_t *s_chime;

typedef enum { JOB_GONG, JOB_REPLY, JOB_CHIME } job_t;
static QueueHandle_t s_jobs;

/* --- gong: see docs/architecture.md "Sound signature" --- */

static const struct {
    float ratio, amp, tau;
} PARTIALS[] = {
    { 1.000f, 0.30f, 2.6f },
    { 1.004f, 0.22f, 2.2f },
    { 1.590f, 0.22f, 1.4f },
    { 2.140f, 0.16f, 0.9f },
    { 2.760f, 0.10f, 0.55f },
    { 3.570f, 0.05f, 0.35f },
    { 4.800f, 0.03f, 0.22f },
};
static const struct {
    int delay_ms;
    float gain;
} ECHOES[] = { { 95, 0.45f }, { 210, 0.30f }, { 370, 0.18f } };

static void synth_gong(void)
{
    float *dry = heap_caps_calloc(GONG_SAMPLES, sizeof(float), MALLOC_CAP_SPIRAM);
    s_gong = heap_caps_malloc(GONG_SAMPLES * sizeof(int16_t), MALLOC_CAP_SPIRAM);
    if (dry == NULL || s_gong == NULL) {
        free(dry);
        return;
    }
    for (size_t p = 0; p < sizeof(PARTIALS) / sizeof(PARTIALS[0]); p++) {
        float step = 2.0f * (float)M_PI * 220.0f * PARTIALS[p].ratio / RATE;
        float decay = expf(-1.0f / (PARTIALS[p].tau * RATE));
        float env = PARTIALS[p].amp, phase = 0.0f;
        for (int i = 0; i < GONG_SAMPLES; i++) {
            phase += step;
            env *= decay;
            dry[i] += sinf(phase) * env;
            if ((i & 0x3FFF) == 0) {
                vTaskDelay(1); /* keep IDLE0 fed; synth is not latency-critical */
            }
        }
    }
    const int attack = RATE * 45 / 1000;
    for (int i = 0; i < attack; i++) {
        dry[i] *= (float)i / attack;
    }
    float peak = 1e-6f;
    for (int i = 0; i < GONG_SAMPLES; i++) {
        float s = dry[i];
        for (size_t e = 0; e < sizeof(ECHOES) / sizeof(ECHOES[0]); e++) {
            int j = i - RATE * ECHOES[e].delay_ms / 1000;
            if (j >= 0) {
                s += dry[j] * ECHOES[e].gain;
            }
        }
        const int fade = RATE * 2 / 5;
        if (i > GONG_SAMPLES - fade) {
            s *= (float)(GONG_SAMPLES - i) / fade;
        }
        dry[i] = s;
        if (fabsf(s) > peak) {
            peak = fabsf(s);
        }
        if ((i & 0x3FFF) == 0) {
            vTaskDelay(1);
        }
    }
    for (int i = 0; i < GONG_SAMPLES; i++) {
        s_gong[i] = (int16_t)(dry[i] * (GONG_PEAK / peak));
    }
    free(dry);
}

/* Short rising two-tone chime played when the wake word fires (you're often
 * not looking at the face, so acknowledge audibly). ~180 ms. */
#define CHIME_SAMPLES (RATE / 5)
static void synth_chime(void)
{
    s_chime = heap_caps_malloc(CHIME_SAMPLES * sizeof(int16_t), MALLOC_CAP_SPIRAM);
    if (s_chime == NULL) {
        return;
    }
    float phase = 0.0f;
    for (int i = 0; i < CHIME_SAMPLES; i++) {
        float freq = (i < CHIME_SAMPLES / 2) ? 660.0f : 990.0f;
        phase += 2.0f * (float)M_PI * freq / RATE;
        float env = 1.0f - (float)i / CHIME_SAMPLES;
        s_chime[i] = (int16_t)(sinf(phase) * 9000.0f * env);
    }
}

/* --- playback worker --- */

static void audio_task(void *arg)
{
    (void)arg;
    synth_gong();
    synth_chime();
    ESP_LOGI(TAG, "gong+chime ready");
    job_t job;
    while (xQueueReceive(s_jobs, &job, portMAX_DELAY) == pdTRUE) {
        /* s_playing gates the mic path: while our own speaker is active, the
         * detect_task neither wakes nor streams (no AEC), so we don't hear
         * ourselves. The tail delay covers the DMA that plays after write()
         * returns, preventing a playback-boundary false wake. */
        if (job == JOB_GONG && s_gong != NULL) {
            esp_codec_dev_write(s_spk, s_gong, GONG_SAMPLES * sizeof(int16_t));
        } else if (job == JOB_CHIME && s_chime != NULL) {
            s_playing = true;
            esp_codec_dev_write(s_spk, s_chime, CHIME_SAMPLES * sizeof(int16_t));
            vTaskDelay(pdMS_TO_TICKS(120));
            s_playing = false;
        } else if (job == JOB_REPLY) {
            s_playing = true;
            esp_codec_dev_write(s_spk, s_reply, s_reply_len);
            vTaskDelay(pdMS_TO_TICKS(250));
            s_playing = false;
            app_on_playback_done();
        }
    }
}

/* AFE wake-word + VAD pipeline. Two tasks:
 *  - feed_task: the SOLE mic reader. Reads mic -> afe->feed(). Always running.
 *  - detect_task: afe->fetch() -> when ARMED, watch for the wake word; when
 *    streaming an utterance, forward AFE-cleaned audio upstream and watch VAD
 *    for end-of-speech. One mic reader ever — no two-readers-on-one-mic hazard. */

static bool afe_init(void)
{
    srmodel_list_t *models = esp_srmodel_init("model");
    char *wn = models ? esp_srmodel_filter(models, ESP_WN_PREFIX, NULL) : NULL;

    afe_config_t *cfg = afe_config_init("M", models, AFE_TYPE_SR, AFE_MODE_LOW_COST);
    if (cfg == NULL) {
        ESP_LOGE(TAG, "afe_config_init failed");
        return false;
    }
    cfg->aec_init = false;   /* no echo canceller in the always-on path (barge-in is later) */
    cfg->se_init = false;    /* single logical mic, no beamforming */
    cfg->ns_init = false;
    cfg->agc_init = false;
    cfg->vad_init = true;    /* VAD gives us hands-free end-of-utterance */
    cfg->vad_mode = VAD_MODE_1;
    cfg->vad_min_speech_ms = 128;
    cfg->vad_min_noise_ms = 600;
    cfg->wakenet_init = (wn != NULL);
    cfg->wakenet_model_name = wn;
    cfg->memory_alloc_mode = AFE_MEMORY_ALLOC_MORE_PSRAM;

    s_afe = esp_afe_handle_from_config(cfg);
    s_afe_data = s_afe ? s_afe->create_from_config(cfg) : NULL;
    if (s_afe == NULL || s_afe_data == NULL) {
        ESP_LOGE(TAG, "AFE create failed");
        return false;
    }
    ESP_LOGI(TAG, "AFE up, wake word = %s", wn ? wn : "NONE");
    return true;
}

static void feed_task(void *arg)
{
    (void)arg;
    int nch = s_afe->get_feed_channel_num(s_afe_data);
    int nsamp = s_afe->get_feed_chunksize(s_afe_data);
    int bytes = nsamp * nch * sizeof(int16_t);
    int16_t *buf = heap_caps_malloc(bytes, MALLOC_CAP_DEFAULT);
    if (buf == NULL) {
        ESP_LOGE(TAG, "feed buffer alloc failed");
        vTaskDelete(NULL);
    }
    for (;;) {
        if (esp_codec_dev_read(s_mic, buf, bytes) == ESP_CODEC_DEV_OK) {
            s_afe->feed(s_afe_data, buf);
        }
    }
}

static void detect_task(void *arg)
{
    (void)arg;
    bool was_streaming = false;
    bool had_speech = false;
    int silence_ms = 0;
    for (;;) {
        afe_fetch_result_t *res = s_afe->fetch(s_afe_data);
        if (res == NULL || res->ret_value == ESP_FAIL) {
            continue;
        }
        int frame_ms = res->data_size / (int)sizeof(int16_t) / (RATE / 1000);

        /* reset VAD accounting whenever streaming (re)starts — from a wake word
         * OR a tap — so a stale silence count can't end the next utterance early */
        if (s_streaming && !was_streaming) {
            had_speech = false;
            silence_ms = 0;
        }
        was_streaming = s_streaming;

        if (!s_streaming) {
            /* ARMED: listen for the wake word (not while our speaker is active) */
            if (!s_playing && res->wakeup_state == WAKENET_DETECTED) {
                app_on_wake();
            }
        } else if (!s_playing) {
            /* UTTERANCE: forward cleaned audio, end on post-speech silence.
             * Skipped while s_playing so the chime/reply doesn't self-stream. */
            gw_send_bin((const uint8_t *)res->data, res->data_size);
            if (res->vad_state == VAD_SPEECH) {
                had_speech = true;
                silence_ms = 0;
            } else {
                silence_ms += frame_ms;
            }
            if (had_speech && silence_ms >= VAD_END_SILENCE_MS) {
                app_on_speech_end();
            }
        }
    }
}

void audio_init(void)
{
    esp_codec_dev_sample_info_t fs = {
        .sample_rate = RATE,
        .channel = 1,
        .bits_per_sample = 16,
    };
    s_spk = bsp_audio_codec_speaker_init();
    if (s_spk != NULL) {
        esp_codec_dev_set_out_vol(s_spk, GONG_VOLUME);
        esp_codec_dev_open(s_spk, &fs);
    }
    s_mic = bsp_audio_codec_microphone_init();
    if (s_mic != NULL) {
        esp_codec_dev_set_in_gain(s_mic, 30.0f);
        esp_codec_dev_open(s_mic, &fs);
    }
    s_reply = heap_caps_malloc(REPLY_MAX, MALLOC_CAP_SPIRAM);
    s_jobs = xQueueCreate(4, sizeof(job_t));
    xTaskCreate(audio_task, "audio", 4096, NULL, 5, NULL);

    if (s_mic != NULL && afe_init()) {
        xTaskCreate(feed_task, "afe_feed", 4096, NULL, 6, NULL);
        xTaskCreate(detect_task, "afe_detect", 8192, NULL, 5, NULL);
    }
    ESP_LOGI(TAG, "audio up (spk=%d mic=%d)", s_spk != NULL, s_mic != NULL);
}

void audio_play_chime(void)
{
    job_t job = JOB_CHIME;
    xQueueSend(s_jobs, &job, 0);
}

void audio_play_gong(void)
{
    job_t job = JOB_GONG;
    xQueueSend(s_jobs, &job, 0);
}

/* Streaming = the detect_task forwards AFE audio upstream. Wake word and touch
 * both begin it; VAD silence and touch both end it. */
void audio_capture_start(void)
{
    s_streaming = true;
}

void audio_capture_stop(void)
{
    s_streaming = false;
}

bool audio_capture_active(void)
{
    return s_streaming;
}

void audio_playback_begin(void)
{
    s_reply_len = 0;
}

void audio_playback_feed(const uint8_t *data, size_t len)
{
    if (s_reply == NULL || s_reply_len + len > REPLY_MAX) {
        return;
    }
    memcpy(s_reply + s_reply_len, data, len);
    s_reply_len += len;
}

void audio_playback_end(void)
{
    job_t job = JOB_REPLY;
    xQueueSend(s_jobs, &job, 0);
}

bool audio_is_playing(void)
{
    return s_playing;
}