#include "voice_interaction_internal.h"

#include <stdlib.h>
#include <string.h>

#define cJSON_malloc voice_malloc_prefer_psram
#define cJSON_realloc voice_realloc_prefer_psram
#include "cJSON.h"
#undef cJSON_realloc
#undef cJSON_malloc
#include "esp_log.h"

static const char *TAG = "voice_interaction";

static void *voice_cjson_malloc_hook(size_t size) {
    return voice_malloc_prefer_psram(size);
}

static void voice_cjson_init_hooks_once(void) {
    static bool s_cjson_hooks_ready;
    if (s_cjson_hooks_ready) {
        return;
    }

    cJSON_Hooks hooks = {
        .malloc_fn = voice_cjson_malloc_hook,
        .free_fn = free,
    };
    cJSON_InitHooks(&hooks);
    s_cjson_hooks_ready = true;
}

void voice_interaction_ws_init_cjson_hooks(void) {
    voice_cjson_init_hooks_once();
}

static char *voice_build_directive_message(const char *action,
                                           const char *directive,
                                           bool with_dialog_id) {
    voice_cjson_init_hooks_once();

    cJSON *root = cJSON_CreateObject();
    if (root == NULL) {
        return NULL;
    }

    cJSON *header = cJSON_AddObjectToObject(root, "header");
    cJSON_AddStringToObject(header, "action", action);
    cJSON_AddStringToObject(header, "task_id", s_voice.task_id);
    cJSON_AddStringToObject(header, "streaming", VOICE_STREAMING_MODE);

    cJSON *payload = cJSON_AddObjectToObject(root, "payload");
    cJSON *input = cJSON_AddObjectToObject(payload, "input");
    cJSON_AddStringToObject(input, "directive", directive);
    if (with_dialog_id && s_voice.dialog_id[0] != '\0') {
        cJSON_AddStringToObject(input, "dialog_id", s_voice.dialog_id);
    }

    char *text = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);
    return text;
}

static char *voice_build_start_message(void) {
    voice_cjson_init_hooks_once();

    cJSON *root = cJSON_CreateObject();
    if (root == NULL) {
        return NULL;
    }

    cJSON *header = cJSON_AddObjectToObject(root, "header");
    cJSON_AddStringToObject(header, "action", "run-task");
    cJSON_AddStringToObject(header, "task_id", s_voice.task_id);
    cJSON_AddStringToObject(header, "streaming", VOICE_STREAMING_MODE);

    cJSON *payload = cJSON_AddObjectToObject(root, "payload");
    cJSON_AddStringToObject(payload, "task_group", VOICE_TASK_GROUP);
    cJSON_AddStringToObject(payload, "task", VOICE_TASK_NAME);
    cJSON_AddStringToObject(payload, "function", VOICE_FUNCTION_NAME);
    cJSON_AddStringToObject(payload, "model", VOICE_MODEL_NAME);

    cJSON *input = cJSON_AddObjectToObject(payload, "input");
    cJSON_AddStringToObject(input, "directive", "Start");
    cJSON_AddStringToObject(input, "workspace_id", CONFIG_TQ_VOICE_WORKSPACE_ID);
    cJSON_AddStringToObject(input, "app_id", CONFIG_TQ_VOICE_APP_ID);

    cJSON *parameters = cJSON_AddObjectToObject(payload, "parameters");
    cJSON *upstream = cJSON_AddObjectToObject(parameters, "upstream");
    cJSON_AddStringToObject(upstream, "type", "AudioOnly");
    cJSON_AddStringToObject(upstream, "mode", "push2talk");
    cJSON_AddStringToObject(upstream, "audio_format", "raw-opus");
    cJSON_AddNumberToObject(upstream, "sample_rate", CONFIG_TQ_VOICE_SAMPLE_RATE);

    cJSON *downstream = cJSON_AddObjectToObject(parameters, "downstream");
    cJSON_AddNumberToObject(downstream, "sample_rate", CONFIG_TQ_VOICE_SAMPLE_RATE);
    cJSON_AddStringToObject(downstream, "audio_format", "raw-opus");
    cJSON_AddNumberToObject(downstream, "frame_size", CONFIG_TQ_VOICE_OPUS_FRAME_MS);
    cJSON_AddNumberToObject(downstream,
                            "bit_rate",
                            CONFIG_TQ_VOICE_OPUS_BITRATE_KBPS);
    if (strlen(CONFIG_TQ_VOICE_TTS_VOICE) > 0) {
        cJSON_AddStringToObject(downstream, "voice", CONFIG_TQ_VOICE_TTS_VOICE);
    }

    cJSON *client_info = cJSON_AddObjectToObject(parameters, "client_info");
    cJSON_AddStringToObject(client_info, "user_id", voice_safe_user_id());
    cJSON *device = cJSON_AddObjectToObject(client_info, "device");
    cJSON_AddStringToObject(device, "uuid", s_voice.device_uuid);

    char *text = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);
    return text;
}

static esp_err_t voice_send_text_immediate(const char *text, int len) {
    if (text == NULL || len <= 0) {
        return ESP_ERR_INVALID_ARG;
    }

    esp_websocket_client_handle_t ws = NULL;
    if (!voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        return ESP_ERR_TIMEOUT;
    }
    ws = s_voice.ws;
    bool connected = s_voice.ws_connected;
    voice_unlock();

    bool transport_connected = (ws != NULL) && esp_websocket_client_is_connected(ws);
    if (connected && !transport_connected && voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        if (s_voice.ws_connected) {
            s_voice.ws_connected = false;
            s_voice.started = false;
            s_voice.tap_active = false;
            s_voice.speech_active = false;
            s_voice.dialog_state = VOICE_DIALOG_STATE_IDLE;
            s_voice.last_event_ms = voice_now_ms();
        }
        voice_unlock();
    }

    if (ws == NULL || !connected || !transport_connected) {
        return ESP_ERR_INVALID_STATE;
    }

    int ret = esp_websocket_client_send_text(ws,
                                             text,
                                             len,
                                             pdMS_TO_TICKS(VOICE_WS_SEND_TIMEOUT_MS));
    if (ret < 0) {
        ESP_LOGW(TAG, "ws_send_text failed: len=%d", len);
        return ESP_FAIL;
    }
    return ESP_OK;
}

static esp_err_t voice_send_text(const char *text, int len) {
    if (text == NULL || len <= 0) {
        return ESP_ERR_INVALID_ARG;
    }

    QueueHandle_t queue = NULL;
    TaskHandle_t ws_text_task = NULL;
    bool session_active = false;
    if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        queue = s_voice.ws_text_queue;
        ws_text_task = s_voice.ws_text_task;
        session_active = s_voice.session_active;
        voice_unlock();
    }

    if (queue == NULL || ws_text_task == NULL || !session_active) {
        return voice_send_text_immediate(text, len);
    }

    char *copy = (char *)voice_malloc_prefer_psram((size_t)len + 1);
    if (copy == NULL) {
        return ESP_ERR_NO_MEM;
    }
    memcpy(copy, text, (size_t)len);
    copy[len] = '\0';

    voice_ws_text_msg_t msg = {
        .text = copy,
        .len = len,
    };
    if (xQueueSend(queue, &msg, pdMS_TO_TICKS(VOICE_WS_SEND_TIMEOUT_MS)) != pdTRUE) {
        free(copy);
        return ESP_ERR_TIMEOUT;
    }
    return ESP_OK;
}

esp_err_t voice_send_directive(const char *action,
                               const char *directive,
                               bool with_dialog_id) {
    char *text = NULL;
    if (!voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        return ESP_ERR_TIMEOUT;
    }
    text = voice_build_directive_message(action, directive, with_dialog_id);
    voice_unlock();

    if (text == NULL) {
        return ESP_ERR_NO_MEM;
    }

    esp_err_t err = voice_send_text(text, (int)strlen(text));
    cJSON_free(text);
    return err;
}

void voice_ws_text_task(void *arg) {
    (void)arg;

    while (true) {
        bool shutdown = false;
        QueueHandle_t queue = NULL;
        if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
            shutdown = s_voice.ws_text_shutdown;
            queue = s_voice.ws_text_queue;
            voice_unlock();
        }

        if (shutdown || queue == NULL) {
            break;
        }

        voice_ws_text_msg_t msg = {0};
        if (xQueueReceive(queue, &msg, pdMS_TO_TICKS(VOICE_WS_TEXT_QUEUE_RX_TIMEOUT_MS)) != pdTRUE) {
            continue;
        }

        if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
            s_voice.ws_text_busy = true;
            voice_unlock();
        }

        esp_err_t err = voice_send_text_immediate(msg.text, msg.len);
        if (err != ESP_OK && voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
            voice_set_last_error_locked("ws text send failed");
            voice_unlock();
        }

        if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
            s_voice.ws_text_busy = false;
            voice_unlock();
        }
        free(msg.text);
    }

    QueueHandle_t queue = NULL;
    if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        queue = s_voice.ws_text_queue;
        s_voice.ws_text_busy = false;
        voice_unlock();
    }
    if (queue != NULL) {
        voice_ws_text_msg_t msg = {0};
        while (xQueueReceive(queue, &msg, 0) == pdTRUE) {
            free(msg.text);
        }
    }

    bool delete_with_caps = false;
    if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
        delete_with_caps = s_voice.ws_text_task_with_caps;
        s_voice.ws_text_task = NULL;
        s_voice.ws_text_task_with_caps = false;
        voice_unlock();
    }
    voice_delete_self_task(delete_with_caps);
}

static void voice_handle_text_chunk(esp_websocket_event_data_t *data) {
    if (data->payload_offset == 0) {
        free(s_voice.text_agg);
        s_voice.text_agg = NULL;
        s_voice.text_agg_size = 0;

        if (data->payload_len <= 0 || data->payload_len > 4096) {
            return;
        }

        s_voice.text_agg = (uint8_t *)voice_calloc_prefer_psram(1, (size_t)data->payload_len + 1);
        if (s_voice.text_agg == NULL) {
            ESP_LOGW(TAG, "text_agg_alloc_failed: len=%d", data->payload_len);
            return;
        }
        s_voice.text_agg_size = (size_t)data->payload_len;
    }

    if (s_voice.text_agg == NULL || s_voice.text_agg_size == 0) {
        return;
    }

    if ((size_t)data->payload_offset + (size_t)data->data_len > s_voice.text_agg_size) {
        free(s_voice.text_agg);
        s_voice.text_agg = NULL;
        s_voice.text_agg_size = 0;
        return;
    }

    memcpy(s_voice.text_agg + data->payload_offset, data->data_ptr, (size_t)data->data_len);

    bool complete = data->fin &&
                    ((size_t)data->payload_offset + (size_t)data->data_len == s_voice.text_agg_size);
    if (!complete) {
        return;
    }

    voice_ws_protocol_handle_text_message((const char *)s_voice.text_agg, s_voice.text_agg_size);
    free(s_voice.text_agg);
    s_voice.text_agg = NULL;
    s_voice.text_agg_size = 0;
}

static void voice_handle_binary_chunk(esp_websocket_event_data_t *data) {
    if (data->payload_offset == 0) {
        free(s_voice.bin_agg);
        s_voice.bin_agg = NULL;
        s_voice.bin_agg_size = 0;

        if (data->payload_len <= 0 || data->payload_len > 16384) {
            return;
        }

        s_voice.bin_agg = (uint8_t *)voice_malloc_prefer_psram((size_t)data->payload_len);
        if (s_voice.bin_agg == NULL) {
            ESP_LOGW(TAG, "bin_agg_alloc_failed: len=%d", data->payload_len);
            return;
        }
        s_voice.bin_agg_size = (size_t)data->payload_len;
    }

    if (s_voice.bin_agg == NULL || s_voice.bin_agg_size == 0) {
        return;
    }

    if ((size_t)data->payload_offset + (size_t)data->data_len > s_voice.bin_agg_size) {
        free(s_voice.bin_agg);
        s_voice.bin_agg = NULL;
        s_voice.bin_agg_size = 0;
        return;
    }

    memcpy(s_voice.bin_agg + data->payload_offset, data->data_ptr, (size_t)data->data_len);

    bool complete = data->fin &&
                    ((size_t)data->payload_offset + (size_t)data->data_len == s_voice.bin_agg_size);
    if (!complete) {
        return;
    }

    voice_ws_audio_handle_downstream_packet(s_voice.bin_agg, s_voice.bin_agg_size);
    free(s_voice.bin_agg);
    s_voice.bin_agg = NULL;
    s_voice.bin_agg_size = 0;
}

void voice_websocket_event_handler(void *handler_args,
                                   esp_event_base_t base,
                                   int32_t event_id,
                                   void *event_data) {
    (void)handler_args;
    (void)base;

    esp_websocket_event_data_t *data = (esp_websocket_event_data_t *)event_data;
    switch ((esp_websocket_event_id_t)event_id) {
        case WEBSOCKET_EVENT_CONNECTED: {
            if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
                s_voice.ws_connected = true;
                s_voice.last_event_ms = voice_now_ms();
                voice_unlock();
            }

            char *start_msg = NULL;
            if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
                start_msg = voice_build_start_message();
                voice_unlock();
            }
            if (start_msg != NULL) {
                (void)voice_send_text(start_msg, (int)strlen(start_msg));
                cJSON_free(start_msg);
            } else {
                ESP_LOGE(TAG, "start_message_build_failed");
            }
            break;
        }
        case WEBSOCKET_EVENT_DISCONNECTED:
            ESP_LOGW(TAG, "ws_disconnected");
            if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
                s_voice.ws_connected = false;
                s_voice.started = false;
                s_voice.tap_active = false;
                s_voice.speech_active = false;
                s_voice.dialog_state = VOICE_DIALOG_STATE_IDLE;
                s_voice.last_event_ms = voice_now_ms();
                voice_unlock();
            }
            voice_close_audio_with_drain(VOICE_AUDIO_DRAIN_WAIT_MS, "ws_disconnected_close");
            break;
        case WEBSOCKET_EVENT_DATA:
            if (data == NULL || data->data_ptr == NULL || data->data_len <= 0) {
                break;
            }
            if (data->op_code == WS_TRANSPORT_OPCODES_TEXT) {
                voice_handle_text_chunk(data);
            } else if (data->op_code == WS_TRANSPORT_OPCODES_BINARY) {
                voice_handle_binary_chunk(data);
            }
            break;
        case WEBSOCKET_EVENT_ERROR:
            ESP_LOGW(TAG, "ws_error");
            if (voice_lock(VOICE_STATUS_LOCK_TIMEOUT_MS)) {
                voice_set_last_error_locked("websocket error");
                s_voice.last_event_ms = voice_now_ms();
                voice_unlock();
            }
            break;
        default:
            break;
    }
}