/* Live Captions: stream mic PCM to Mac mini and save final text on the SD card. */
#include <tt_app.h>
#include <tt_lvgl.h>
#include <tt_lvgl_toolbar.h>
#include <tactility/device.h>
#include <tactility/drivers/audio_stream.h>

#include <cJSON.h>
#include <esp_log.h>
#include <esp_random.h>
#include <freertos/FreeRTOS.h>
#include <freertos/semphr.h>
#include <freertos/task.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/stat.h>

#include "websocket.h"

/* These legacy names are exported by the flashed 0.8.0-dev firmware. */
struct Device* device_find_by_name(const char* name);
struct Device* device_find_first_by_type(const struct DeviceType* type);

#define TAG "LiveCaptions"
#define DEFAULT_ENDPOINT "ws://192.168.68.102:8645/api/esp32/captions/ws"
#define DEFAULT_DEVICE_ID "tactility-14c19d1a790"
#define PCM_BUFFER_BYTES 1024U
#define EVENT_BUFFER_BYTES 4096U
#define DISPLAY_WORD_LIMIT 50

typedef enum {
    CAPTION_IDLE,
    CAPTION_CONNECTING,
    CAPTION_LISTENING,
    CAPTION_PROCESSING,
    CAPTION_FAILED,
} CaptionState;

typedef struct {
    AppHandle app;
    volatile bool visible;
    volatile bool capture_audio;
    volatile bool stop_requested;
    volatile bool session_active;
    volatile bool socket_failed;
    int fd;
    CaptionState state;
    char endpoint[128];
    char device_id[64];
    char api_key[128];
    char detail[96];
    char caption[768];
    char last_final[768];
    struct Device* stream_dev;
    AudioStreamHandle input_handle;
    TaskHandle_t worker;
    TaskHandle_t receiver;
    SemaphoreHandle_t socket_lock;
    SemaphoreHandle_t audio_lock;
    lv_obj_t* caption_label;
    lv_obj_t* status_label;
} CaptionContext;

static void update_ui(CaptionContext* ctx) {
    if (!ctx->visible || !tt_lvgl_lock(pdMS_TO_TICKS(100))) return;
    lv_label_set_text(ctx->caption_label, ctx->caption);
    const char* status = ctx->state == CAPTION_LISTENING ? "Connected" : (ctx->state == CAPTION_CONNECTING ? "Connecting" : (ctx->state == CAPTION_FAILED ? "Failed — Reconnecting" : "Connecting"));
    lv_label_set_text(ctx->status_label, status);
    tt_lvgl_unlock();
}

static void set_state(CaptionContext* ctx, CaptionState state, const char* detail) {
    ctx->state = state;
    snprintf(ctx->detail, sizeof(ctx->detail), "%s", detail ? detail : "");
    update_ui(ctx);
}

static bool find_audio_stream_device(CaptionContext* ctx) {
    ctx->stream_dev = device_find_by_name("audio-stream");
    if (ctx->stream_dev == NULL) ctx->stream_dev = device_find_first_by_type(&AUDIO_STREAM_TYPE);
    return ctx->stream_dev != NULL;
}

static bool open_input_stream(CaptionContext* ctx) {
    if (ctx->stream_dev == NULL) return false;
    if (ctx->input_handle != NULL) return true;
    struct AudioStreamConfig config = {.sample_rate = 16000, .bits_per_sample = 16, .channels = 1};
    if (audio_stream_open_input(ctx->stream_dev, &config, &ctx->input_handle) != ERROR_NONE) return false;
    audio_stream_set_mute(ctx->stream_dev, AUDIO_CODEC_DIR_INPUT, false);
    audio_stream_set_volume(ctx->stream_dev, AUDIO_CODEC_DIR_INPUT, 100.0f);
    return true;
}

static void close_input_stream(CaptionContext* ctx) {
    if (ctx->input_handle != NULL) {
        audio_stream_close(ctx->input_handle);
        ctx->input_handle = NULL;
    }
}

static int send_locked(CaptionContext* ctx, const uint8_t* data, size_t length, bool binary) {
    if (ctx->fd < 0 || xSemaphoreTake(ctx->socket_lock, pdMS_TO_TICKS(500)) != pdTRUE) return -1;
    int result = ws_send(ctx->fd, data, length, binary);
    xSemaphoreGive(ctx->socket_lock);
    return result;
}

static bool parse_endpoint(const char* url, char* host, size_t host_size, int* port, char* path, size_t path_size) {
    if (url == NULL || strncmp(url, "ws://", 5) != 0) return false;
    const char* authority = url + 5;
    const char* slash = strchr(authority, '/');
    const char* end = slash ? slash : authority + strlen(authority);
    const char* colon = NULL;
    for (const char* p = authority; p < end; ++p) if (*p == ':') colon = p;
    size_t host_len = (size_t)((colon ? colon : end) - authority);
    if (host_len == 0 || host_len >= host_size) return false;
    memcpy(host, authority, host_len); host[host_len] = '\0';
    *port = 80;
    if (colon != NULL) {
        *port = atoi(colon + 1);
        if (*port < 1 || *port > 65535) return false;
    }
    const char* wire_path = slash ? slash : "/";
    if (strlen(wire_path) >= path_size) return false;
    snprintf(path, path_size, "%s", wire_path);
    return true;
}

static void append_final_caption(const char* text) {
    if (text == NULL || !*text) return;
    mkdir("/sdcard/captions", 0755);
    time_t now = time(NULL);
    struct tm local;
    localtime_r(&now, &local);
    char filename[32];
    strftime(filename, sizeof(filename), "%Y-%m-%d.txt", &local);
    char stamp[16];
    strftime(stamp, sizeof(stamp), "%H:%M:%S", &local);
    char path[128];
    snprintf(path, sizeof(path), "/sdcard/captions/%s", filename);
    FILE* log = fopen(path, "a");
    if (log == NULL) { ESP_LOGE(TAG, "could not append %s", path); return; }
    fprintf(log, "[%s] %s\n", stamp, text);
    fclose(log);
}

static bool is_space_char(char value) { return value == ' ' || value == '\n' || value == '\r' || value == '	'; }

static void copy_recent_words(char* output, size_t output_size, const char* text) {
    const char* starts[DISPLAY_WORD_LIMIT];
    int count = 0;
    bool in_word = false;
    for (const char* cursor = text; *cursor; ++cursor) {
        if (is_space_char(*cursor)) { in_word = false; continue; }
        if (!in_word) { starts[count % DISPLAY_WORD_LIMIT] = cursor; ++count; in_word = true; }
    }
    const char* first = count > DISPLAY_WORD_LIMIT ? starts[count % DISPLAY_WORD_LIMIT] : text;
    snprintf(output, output_size, "%s", first);
}

static void display_caption(CaptionContext* ctx, const char* text, bool final) {
    if (text == NULL || !*text) return;
    copy_recent_words(ctx->caption, sizeof(ctx->caption), text);
    if (final && strcmp(ctx->last_final, text) != 0) {
        snprintf(ctx->last_final, sizeof(ctx->last_final), "%s", text);
        append_final_caption(text);
        ctx->state = CAPTION_IDLE;
    }
    update_ui(ctx);
}

static void handle_event(CaptionContext* ctx, const char* json) {
    cJSON* root = cJSON_Parse(json);
    if (root == NULL) return;
    cJSON* event = cJSON_GetObjectItem(root, "event");
    cJSON* text = cJSON_GetObjectItem(root, "text");
    if (!cJSON_IsString(event)) { cJSON_Delete(root); return; }
    const char* name = event->valuestring;
    if (strcmp(name, "ready") == 0) {
        set_state(ctx, CAPTION_CONNECTING, "Starting caption stream");
    } else if (strcmp(name, "listening") == 0) {
        set_state(ctx, CAPTION_LISTENING, "Listening — press Stop when finished");
    } else if (strcmp(name, "state") == 0) {
        cJSON* remote_state = cJSON_GetObjectItem(root, "state");
        if (cJSON_IsString(remote_state) && strcmp(remote_state->valuestring, "listening") == 0) {
            set_state(ctx, CAPTION_LISTENING, "Listening");
        } else if (cJSON_IsString(remote_state) && strcmp(remote_state->valuestring, "processing") == 0) {
            set_state(ctx, CAPTION_PROCESSING, "Captioning…");
        }
    } else if (strcmp(name, "draft") == 0 || strcmp(name, "interim_transcript") == 0 || strcmp(name, "review") == 0) {
        if (cJSON_IsString(text)) display_caption(ctx, text->valuestring, false);
    } else if (strcmp(name, "transcript") == 0 || strcmp(name, "final") == 0) {
        cJSON* is_final = cJSON_GetObjectItem(root, "isFinal");
        if (cJSON_IsString(text) && (!cJSON_IsBool(is_final) || cJSON_IsTrue(is_final) || strcmp(name, "final") == 0)) {
            display_caption(ctx, text->valuestring, true);
        }
    } else if (strcmp(name, "thinking") == 0) {
        set_state(ctx, CAPTION_PROCESSING, "Final captioning…");
    } else if (strcmp(name, "error") == 0) {
        ctx->socket_failed = true;
        set_state(ctx, CAPTION_FAILED, "Fail to connect");
    }
    cJSON_Delete(root);
}

static void receiver_task(void* argument) {
    CaptionContext* ctx = argument;
    uint8_t* buffer = malloc(EVENT_BUFFER_BYTES + 1U);
    if (buffer == NULL) { ctx->socket_failed = true; ctx->receiver = NULL; vTaskDelete(NULL); }
    while (ctx->visible && ctx->fd >= 0) {
        int opcode = 0; bool complete = false;
        int received = ws_recv(ctx->fd, &opcode, &complete, buffer, EVENT_BUFFER_BYTES);
        if (received < 0 || !complete) break;
        if (opcode == 0x01) { buffer[received] = '\0'; handle_event(ctx, (const char*)buffer); }
        else if (opcode == 0x09) {
            if (xSemaphoreTake(ctx->socket_lock, pdMS_TO_TICKS(500)) == pdTRUE) {
                ws_send_pong(ctx->fd, buffer, (size_t)received); xSemaphoreGive(ctx->socket_lock);
            }
        } else if (opcode == 0x08) break;
    }
    ctx->session_active = false;
    ctx->receiver = NULL;
    vTaskDelete(NULL);
}

static void load_config(CaptionContext* ctx) {
    snprintf(ctx->endpoint, sizeof(ctx->endpoint), "%s", DEFAULT_ENDPOINT);
    snprintf(ctx->device_id, sizeof(ctx->device_id), "%s", DEFAULT_DEVICE_ID);
    ctx->api_key[0] = '\0';
    char path[256]; size_t size = sizeof(path);
    tt_app_get_user_data_child_path(ctx->app, "config.json", path, &size);
    FILE* file = fopen(path, "r");
    if (file == NULL) return;
    char raw[512]; size_t bytes = fread(raw, 1, sizeof(raw) - 1, file); fclose(file); raw[bytes] = '\0';
    cJSON* root = cJSON_Parse(raw);
    cJSON* endpoint = root ? cJSON_GetObjectItem(root, "server_url") : NULL;
    cJSON* device = root ? cJSON_GetObjectItem(root, "device_id") : NULL;
    cJSON* key = root ? cJSON_GetObjectItem(root, "api_key") : NULL;
    if (cJSON_IsString(endpoint)) snprintf(ctx->endpoint, sizeof(ctx->endpoint), "%s", endpoint->valuestring);
    if (cJSON_IsString(device)) snprintf(ctx->device_id, sizeof(ctx->device_id), "%s", device->valuestring);
    if (cJSON_IsString(key)) snprintf(ctx->api_key, sizeof(ctx->api_key), "%s", key->valuestring);
    cJSON_Delete(root);
}

static void worker_task(void* argument) {
    CaptionContext* ctx = argument;
    char host[64], path[96]; int port = 0;
    if (!parse_endpoint(ctx->endpoint, host, sizeof(host), &port, path, sizeof(path))) {
        set_state(ctx, CAPTION_FAILED, "Fail to connect"); ctx->worker = NULL; vTaskDelete(NULL);
    }
    set_state(ctx, CAPTION_CONNECTING, "Connecting to Mac mini");
    ctx->fd = ws_connect(host, port, path, ctx->device_id, ctx->api_key);
    if (ctx->fd < 0) { set_state(ctx, CAPTION_FAILED, "Fail to connect"); ctx->worker = NULL; vTaskDelete(NULL); }
    char start[256];
    snprintf(start, sizeof(start), "{\"v\":1,\"event\":\"start\",\"session_id\":\"cap-%08lx\",\"device_id\":\"%s\",\"audio\":{\"format\":\"pcm_s16le\",\"sample_rate\":16000,\"channels\":1,\"sample_width\":2}}", (unsigned long)esp_random(), ctx->device_id);
    if (send_locked(ctx, (const uint8_t*)start, strlen(start), false) != 0 || !open_input_stream(ctx)) {
        set_state(ctx, CAPTION_FAILED, "Fail to connect"); ws_close(ctx->fd); ctx->fd = -1; ctx->worker = NULL; vTaskDelete(NULL);
    }
    ctx->session_active = true;
    xTaskCreate(receiver_task, "caption_rx", 6144, ctx, 6, &ctx->receiver);
    set_state(ctx, CAPTION_LISTENING, "Listening — press Stop when finished");
    uint8_t pcm[PCM_BUFFER_BYTES];
    while (ctx->visible && ctx->capture_audio && !ctx->socket_failed) {
        size_t bytes = 0;
        xSemaphoreTake(ctx->audio_lock, portMAX_DELAY);
        error_t result = audio_stream_read(ctx->input_handle, pcm, sizeof(pcm), &bytes, pdMS_TO_TICKS(100));
        xSemaphoreGive(ctx->audio_lock);
        if (result == ERROR_NONE && bytes > 0 && (bytes % 2U) == 0 && send_locked(ctx, pcm, bytes, true) != 0) ctx->socket_failed = true;
    }
    xSemaphoreTake(ctx->audio_lock, portMAX_DELAY); close_input_stream(ctx); xSemaphoreGive(ctx->audio_lock);
    if (ctx->stop_requested && !ctx->socket_failed) {
        const char* stop = "{\"event\":\"stop\"}";
        send_locked(ctx, (const uint8_t*)stop, strlen(stop), false);
        set_state(ctx, CAPTION_PROCESSING, "Final captioning…");
        for (unsigned i = 0; ctx->session_active && i < 150; ++i) vTaskDelay(pdMS_TO_TICKS(100));
    }
    if (ctx->fd >= 0) { ws_send_close(ctx->fd); ws_close(ctx->fd); ctx->fd = -1; }
    bool reconnect = ctx->socket_failed && ctx->visible;
    if (reconnect) set_state(ctx, CAPTION_FAILED, "Reconnecting");
    else if (ctx->state == CAPTION_PROCESSING) set_state(ctx, CAPTION_IDLE, "");
    ctx->worker = NULL;
    if (reconnect) {
        vTaskDelay(pdMS_TO_TICKS(3000));
        if (ctx->visible) { ctx->socket_failed = false; ctx->capture_audio = true; xTaskCreate(worker_task, "caption_tx", 8192, ctx, 5, &ctx->worker); }
    }
    vTaskDelete(NULL);
}

static void* create_data(void) { CaptionContext* ctx = calloc(1, sizeof(*ctx)); if (ctx) ctx->fd = -1; return ctx; }
static void destroy_data(void* data) { free(data); }
static void on_create(AppHandle app, void* data) { ((CaptionContext*)data)->app = app; }

static void on_show(AppHandle app, void* data, lv_obj_t* parent) {
    CaptionContext* ctx = data; ctx->visible = true; load_config(ctx); find_audio_stream_device(ctx);
    ctx->socket_lock = xSemaphoreCreateMutex(); ctx->audio_lock = xSemaphoreCreateMutex();
    lv_obj_t* toolbar = tt_lvgl_toolbar_create_for_app(parent, app);
    lv_obj_align(toolbar, LV_ALIGN_TOP_MID, 0, 0);
    ctx->caption_label = lv_label_create(parent);
    lv_obj_set_width(ctx->caption_label, lv_pct(88));
    lv_label_set_long_mode(ctx->caption_label, LV_LABEL_LONG_WRAP);
    lv_obj_set_style_text_align(ctx->caption_label, LV_TEXT_ALIGN_CENTER, 0);
    lv_obj_align(ctx->caption_label, LV_ALIGN_CENTER, 0, 0);
    ctx->status_label = lv_label_create(parent);
    lv_obj_align(ctx->status_label, LV_ALIGN_BOTTOM_MID, 0, -10);
    ctx->caption[0] = '\0';
    if (ctx->stream_dev != NULL && ctx->socket_lock != NULL && ctx->audio_lock != NULL) {
        ctx->capture_audio = true; ctx->stop_requested = false; ctx->socket_failed = false;
        xTaskCreate(worker_task, "caption_tx", 8192, ctx, 5, &ctx->worker);
    }
}

static void on_hide(AppHandle app, void* data) {
    (void)app; CaptionContext* ctx = data; ctx->visible = false; ctx->capture_audio = false; ctx->stop_requested = false;
    if (ctx->fd >= 0) { ws_send_close(ctx->fd); ws_close(ctx->fd); ctx->fd = -1; }
    for (unsigned i = 0; (ctx->worker || ctx->receiver) && i < 100; ++i) vTaskDelay(pdMS_TO_TICKS(10));
    close_input_stream(ctx);
    if (ctx->socket_lock) { vSemaphoreDelete(ctx->socket_lock); ctx->socket_lock = NULL; }
    if (ctx->audio_lock) { vSemaphoreDelete(ctx->audio_lock); ctx->audio_lock = NULL; }
}

int main(int argc, char* argv[]) {
    (void)argc; (void)argv;
    tt_app_register((AppRegistration){.createData=create_data,.destroyData=destroy_data,.onCreate=on_create,.onShow=on_show,.onHide=on_hide});
    return 0;
}