#include "audio_client.h"
#include "codec_bsp.h"
#include "esp_websocket_client.h"
#include "esp_log.h"
#include "cJSON.h"

#include <string.h>
#include <freertos/FreeRTOS.h>
#include <freertos/queue.h>
#include <freertos/task.h>
#include <esp_heap_caps.h>

static const char *TAG = "audio_client";

#define AUDIO_CHUNK_SIZE    4096
#define PCM_QUEUE_DEPTH     10
#define PLAYBACK_STACK_SIZE (4 * 1024)
#define PLAYBACK_PRIORITY   4
#define WS_BUFFER_SIZE      8192

/* Status image constants */
#define STATUS_IMG_W     120
#define STATUS_IMG_H     120
#define STATUS_IMG_BYTES (STATUS_IMG_W * STATUS_IMG_H / 8)  /* 1800 */

static esp_websocket_client_handle_t s_client = NULL;
static CodecPort *s_codec = NULL;
static QueueHandle_t s_pcm_queue = NULL;
static TaskHandle_t s_playback_task = NULL;
static volatile audio_state_t s_state = AUDIO_IDLE;
static volatile bool s_playing = false;

/* Status image state */
static uint8_t s_img_buf[STATUS_IMG_BYTES];
static lv_img_dsc_t s_img_dsc;
static volatile bool s_img_pending = false;   /* expecting binary frame with image data */
static volatile bool s_img_updated = false;   /* new image ready for UI consumption */
static TaskHandle_t s_img_notify_task = NULL;  /* task to wake on new image */

/* Forward declarations */
static void playback_task(void *arg);
static void ws_event_handler(void *arg, esp_event_base_t event_base,
                             int32_t event_id, void *event_data);

/* ---------- Queue helpers ---------- */

static void flush_queue(void)
{
    uint8_t *chunk;
    while (xQueueReceive(s_pcm_queue, &chunk, 0) == pdTRUE) {
        heap_caps_free(chunk);
    }
}

/* ---------- WebSocket event handler ---------- */

static void handle_text_frame(const char *data, int len)
{
    cJSON *root = cJSON_ParseWithLength(data, len);
    if (!root) {
        ESP_LOGW(TAG, "JSON parse failed");
        return;
    }

    cJSON *type = cJSON_GetObjectItem(root, "type");
    if (!cJSON_IsString(type)) {
        cJSON_Delete(root);
        return;
    }

    if (strcmp(type->valuestring, "status_image") == 0) {
        ESP_LOGI(TAG, "Status image header received");
        s_img_pending = true;
        cJSON_Delete(root);
        return;
    }

    if (strcmp(type->valuestring, "alarm_start") == 0) {
        int sr   = 24000;
        int ch   = 2;
        int bits = 16;

        cJSON *item;
        item = cJSON_GetObjectItem(root, "sample_rate");
        if (cJSON_IsNumber(item)) sr = item->valueint;
        item = cJSON_GetObjectItem(root, "channels");
        if (cJSON_IsNumber(item)) ch = item->valueint;
        item = cJSON_GetObjectItem(root, "bits");
        if (cJSON_IsNumber(item)) bits = item->valueint;

        ESP_LOGI(TAG, "Alarm start: %dHz %dch %dbit", sr, ch, bits);

        /* Flush any stale data */
        flush_queue();

        /* Open codec for playback */
        s_codec->CodecPort_SetInfo("es8311", 1, sr, ch, bits);
        s_codec->CodecPort_SetSpeakerVol(70);

        s_playing = true;
        s_state = AUDIO_PLAYING;

    } else if (strcmp(type->valuestring, "alarm_stop") == 0) {
        ESP_LOGI(TAG, "Alarm stop");
        s_playing = false;

        /* Let playback task drain remaining chunks, then close */
        vTaskDelay(pdMS_TO_TICKS(100));
        flush_queue();
        s_codec->CodecPort_CloseSpeaker();

        s_state = AUDIO_CONNECTED;
    }

    cJSON_Delete(root);
}

static void handle_binary_frame(const uint8_t *data, int len)
{
    /* Status image binary payload */
    if (s_img_pending) {
        if (len == STATUS_IMG_BYTES) {
            memcpy(s_img_buf, data, STATUS_IMG_BYTES);
            s_img_updated = true;
            if (s_img_notify_task) {
                xTaskNotifyGive(s_img_notify_task);
            }
            ESP_LOGI(TAG, "Status image received (%d bytes)", len);
        } else {
            ESP_LOGW(TAG, "Status image size mismatch: got %d, expected %d", len, STATUS_IMG_BYTES);
        }
        s_img_pending = false;
        return;
    }

    if (!s_playing) return;

    uint8_t *chunk = (uint8_t *)heap_caps_malloc(len, MALLOC_CAP_SPIRAM);
    if (!chunk) {
        ESP_LOGW(TAG, "PSRAM alloc failed (%d bytes)", len);
        return;
    }
    memcpy(chunk, data, len);

    if (xQueueSend(s_pcm_queue, &chunk, 0) != pdTRUE) {
        ESP_LOGW(TAG, "PCM queue full, dropping chunk");
        heap_caps_free(chunk);
    }
}

static void ws_event_handler(void *arg, esp_event_base_t event_base,
                             int32_t event_id, void *event_data)
{
    esp_websocket_event_data_t *ev = (esp_websocket_event_data_t *)event_data;

    switch (event_id) {
    case WEBSOCKET_EVENT_CONNECTED:
        ESP_LOGI(TAG, "Audio WS connected");
        s_state = AUDIO_CONNECTED;
        break;

    case WEBSOCKET_EVENT_DISCONNECTED:
        ESP_LOGW(TAG, "Audio WS disconnected");
        s_playing = false;
        flush_queue();
        s_state = AUDIO_IDLE;
        break;

    case WEBSOCKET_EVENT_DATA:
        if (ev->op_code == 0x01 && ev->data_len > 0) {
            handle_text_frame(ev->data_ptr, ev->data_len);
        } else if (ev->op_code == 0x02 && ev->data_len > 0) {
            handle_binary_frame((const uint8_t *)ev->data_ptr, ev->data_len);
        }
        break;

    case WEBSOCKET_EVENT_ERROR:
        ESP_LOGE(TAG, "Audio WS error");
        s_playing = false;
        s_state = AUDIO_ERROR;
        break;

    default:
        break;
    }
}

/* ---------- Playback task ---------- */

static void playback_task(void *arg)
{
    uint8_t *chunk;

    for (;;) {
        if (xQueueReceive(s_pcm_queue, &chunk, pdMS_TO_TICKS(500)) == pdTRUE) {
            if (s_playing) {
                s_codec->CodecPort_PlayWrite(chunk, AUDIO_CHUNK_SIZE);
            }
            heap_caps_free(chunk);
        }
    }
}

/* ---------- Public API ---------- */

void audio_client_init(const char *uri, void *codec)
{
    s_codec = (CodecPort *)codec;

    /* Initialize status image descriptor */
    memset(&s_img_dsc, 0, sizeof(s_img_dsc));
    s_img_dsc.header.cf = LV_IMG_CF_ALPHA_1BIT;
    s_img_dsc.header.w = STATUS_IMG_W;
    s_img_dsc.header.h = STATUS_IMG_H;
    s_img_dsc.data_size = STATUS_IMG_BYTES;
    s_img_dsc.data = s_img_buf;

    s_pcm_queue = xQueueCreate(PCM_QUEUE_DEPTH, sizeof(uint8_t *));
    if (!s_pcm_queue) {
        ESP_LOGE(TAG, "Failed to create PCM queue");
        return;
    }

    esp_websocket_client_config_t config = {};
    config.uri = uri;
    config.reconnect_timeout_ms = 5000;
    config.buffer_size = WS_BUFFER_SIZE;

    s_client = esp_websocket_client_init(&config);
    esp_websocket_register_events(s_client, WEBSOCKET_EVENT_ANY, ws_event_handler, NULL);

    ESP_LOGI(TAG, "Audio client initialized: %s", uri);
}

void audio_client_start(void)
{
    if (!s_client) return;

    /* Create playback task pinned to Core 1 */
    xTaskCreatePinnedToCore(playback_task, "audio_play", PLAYBACK_STACK_SIZE,
                            NULL, PLAYBACK_PRIORITY, &s_playback_task, 1);

    esp_websocket_client_start(s_client);
    ESP_LOGI(TAG, "Audio client started");
}

void audio_client_stop(void)
{
    if (!s_client) return;

    s_playing = false;
    esp_websocket_client_stop(s_client);
    flush_queue();

    if (s_playback_task) {
        vTaskDelete(s_playback_task);
        s_playback_task = NULL;
    }

    s_state = AUDIO_IDLE;
    ESP_LOGI(TAG, "Audio client stopped");
}

audio_state_t audio_client_get_state(void)
{
    return s_state;
}

void audio_client_set_image_notify_task(TaskHandle_t task)
{
    s_img_notify_task = task;
}

const lv_img_dsc_t *audio_client_get_status_image(bool *updated)
{
    if (updated) {
        *updated = s_img_updated;
        if (s_img_updated) {
            s_img_updated = false;
        }
    }
    return &s_img_dsc;
}