#include "infinite_canvas/agent_service.h"

#include "mrjunejune/inference_bridge.h"

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define CANVAS_AGENT_WARM_SESSION_COUNT 4

typedef struct {
    pthread_mutex_t mutex;
    Canvas_Agent_Service *p_service;
    Inference_Bridge *p_bridge;
    uint32 request_sequence;
    uint32 warmed_sessions;
    char request_id[64];
    char turn_content[CANVAS_AGENT_RESPONSE_CAPACITY];
    char turn_error[512];
    char turn_prompt[
        CANVAS_CONTEXT_MAX_LENGTH + CANVAS_ENTITY_TEXT_CAPACITY + 512];
} Canvas_Agent_Copilot;

static boolean Canvas_Agent_Event_Matches(
    const Canvas_Agent_Copilot *p_copilot,
    const Inference_Event *p_event)
{
    return p_event->request_id &&
        strcmp(p_copilot->request_id, p_event->request_id) == 0;
}

static void Canvas_Agent_Handle_Event(
    const Inference_Event *p_event,
    void *p_user_data)
{
    Canvas_Agent_Copilot *p_copilot =
        (Canvas_Agent_Copilot *)p_user_data;
    if (!p_copilot || !p_event) return;

    pthread_mutex_lock(&p_copilot->mutex);
    if (strcmp(p_event->type, "bridge.closed") == 0) {
        if (p_copilot->p_service->status == CANVAS_AGENT_WORKING) {
            snprintf(
                p_copilot->p_service->error,
                sizeof(p_copilot->p_service->error),
                "%s",
                p_event->error_message && p_event->error_message[0]
                    ? p_event->error_message
                    : "Copilot orchestration sidecar closed");
            p_copilot->p_service->status = CANVAS_AGENT_ERROR;
        }
        pthread_mutex_unlock(&p_copilot->mutex);
        return;
    }
    if (strcmp(p_event->type, "session.warmed") == 0 &&
        p_event->request_id &&
        strncmp(p_event->request_id, "canvas-warm-", 12) == 0) {
        p_copilot->warmed_sessions++;
        pthread_mutex_unlock(&p_copilot->mutex);
        return;
    }
    if (!Canvas_Agent_Event_Matches(p_copilot, p_event)) {
        pthread_mutex_unlock(&p_copilot->mutex);
        return;
    }

    if (strcmp(p_event->type, "assistant.completed") == 0) {
        snprintf(
            p_copilot->turn_content,
            sizeof(p_copilot->turn_content),
            "%s",
            p_event->content ? p_event->content : "");
    } else if (strcmp(p_event->type, "turn.error") == 0) {
        snprintf(
            p_copilot->turn_error,
            sizeof(p_copilot->turn_error),
            "%s",
            p_event->error_message && p_event->error_message[0]
                ? p_event->error_message
                : "Copilot orchestration failed");
    } else if (strcmp(p_event->type, "turn.done") == 0) {
        if (p_event->failed || p_copilot->turn_error[0]) {
            snprintf(
                p_copilot->p_service->error,
                sizeof(p_copilot->p_service->error),
                "%s",
                p_copilot->turn_error[0]
                    ? p_copilot->turn_error
                    : "Copilot orchestration failed");
            p_copilot->p_service->status = CANVAS_AGENT_ERROR;
        } else if (p_copilot->turn_content[0]) {
            snprintf(
                p_copilot->p_service->response,
                sizeof(p_copilot->p_service->response),
                "%s",
                p_copilot->turn_content);
            p_copilot->p_service->status = CANVAS_AGENT_READY;
        } else {
            snprintf(
                p_copilot->p_service->error,
                sizeof(p_copilot->p_service->error),
                "Copilot returned an empty response");
            p_copilot->p_service->status = CANVAS_AGENT_ERROR;
        }
    }
    pthread_mutex_unlock(&p_copilot->mutex);
}

boolean Canvas_Agent_Service_Init(
    Canvas_Agent_Service *p_service,
    Dowa_Arena *p_arena)
{
    memset(p_service, 0, sizeof(*p_service));
    p_service->p_arena = p_arena;
    const char *p_sidecar =
        getenv("INFINITE_CANVAS_COPILOT_SIDECAR_PATH");
    const char *p_cli =
        getenv("INFINITE_CANVAS_COPILOT_CLI_PATH");
    if (!p_sidecar || !p_sidecar[0] || !p_cli || !p_cli[0]) {
        return FALSE;
    }

    Canvas_Agent_Copilot *p_copilot = (Canvas_Agent_Copilot *)
        Dowa_Arena_Allocate(p_arena, sizeof(*p_copilot));
    if (!p_copilot) return FALSE;
    memset(p_copilot, 0, sizeof(*p_copilot));
    if (pthread_mutex_init(&p_copilot->mutex, NULL) != 0) {
        return FALSE;
    }
    p_copilot->p_service = p_service;
    p_copilot->p_bridge = Inference_Bridge_Create(
        p_sidecar,
        p_cli,
        Canvas_Agent_Handle_Event,
        p_copilot);
    if (!p_copilot->p_bridge ||
        !Inference_Bridge_Start(p_copilot->p_bridge)) {
        if (p_copilot->p_bridge) {
            Inference_Bridge_Destroy(p_copilot->p_bridge);
        }
        pthread_mutex_destroy(&p_copilot->mutex);
        return FALSE;
    }
    const char *warm_sessions[CANVAS_AGENT_WARM_SESSION_COUNT] = {
        "infinite-canvas-orchestrator",
        "infinite-canvas-worker-1",
        "infinite-canvas-worker-2",
        "infinite-canvas-worker-3",
    };
    for (uint32 index = 0;
        index < CANVAS_AGENT_WARM_SESSION_COUNT;
        index++) {
        char request_id[32];
        snprintf(
            request_id,
            sizeof(request_id),
            "canvas-warm-%u",
            index);
        if (!Inference_Bridge_Warm_Conversation(
                p_copilot->p_bridge,
                request_id,
                warm_sessions[index],
                INFERENCE_PROMPT_PROFILE_CANVAS_ORCHESTRATOR,
                1,
                1,
                index == 0 ? TRUE : FALSE)) {
            Inference_Bridge_Destroy(p_copilot->p_bridge);
            pthread_mutex_destroy(&p_copilot->mutex);
            return FALSE;
        }
    }
    for (int32 attempt = 0; attempt < 2000; attempt++) {
        pthread_mutex_lock(&p_copilot->mutex);
        boolean ready =
            p_copilot->warmed_sessions == CANVAS_AGENT_WARM_SESSION_COUNT;
        pthread_mutex_unlock(&p_copilot->mutex);
        if (ready) break;
        usleep(10000);
    }
    p_service->p_native = p_copilot;
    p_service->status = CANVAS_AGENT_IDLE;
    return TRUE;
}

boolean Canvas_Agent_Service_Submit(
    Canvas_Agent_Service *p_service,
    const char *p_prompt,
    const char *p_context)
{
    Canvas_Agent_Copilot *p_copilot =
        (Canvas_Agent_Copilot *)p_service->p_native;
    if (!p_copilot || !p_prompt || !p_prompt[0]) return FALSE;

    pthread_mutex_lock(&p_copilot->mutex);
    if (p_service->status == CANVAS_AGENT_WORKING) {
        pthread_mutex_unlock(&p_copilot->mutex);
        return FALSE;
    }
    p_copilot->request_sequence++;
    snprintf(
        p_copilot->request_id,
        sizeof(p_copilot->request_id),
        "canvas-%u",
        p_copilot->request_sequence);
    snprintf(
        p_copilot->turn_prompt,
        sizeof(p_copilot->turn_prompt),
        "Visible canvas context:\n%s\n\nNew user thought:\n%s",
        p_context ? p_context : "",
        p_prompt);
    p_copilot->turn_content[0] = '\0';
    p_copilot->turn_error[0] = '\0';
    p_service->response[0] = '\0';
    p_service->error[0] = '\0';
    p_service->status = CANVAS_AGENT_WORKING;
    char request_id[sizeof(p_copilot->request_id)];
    snprintf(request_id, sizeof(request_id), "%s", p_copilot->request_id);
    pthread_mutex_unlock(&p_copilot->mutex);

    if (!Inference_Bridge_Start_Turn(
            p_copilot->p_bridge,
            request_id,
            "infinite-canvas-orchestrator",
            p_copilot->turn_prompt,
            INFERENCE_PROMPT_PROFILE_CANVAS_ORCHESTRATOR,
            1,
            1,
            NULL,
            0)) {
        pthread_mutex_lock(&p_copilot->mutex);
        snprintf(
            p_service->error,
            sizeof(p_service->error),
            "Unable to start the Copilot orchestration turn");
        p_service->status = CANVAS_AGENT_ERROR;
        pthread_mutex_unlock(&p_copilot->mutex);
        return FALSE;
    }
    return TRUE;
}

Canvas_Agent_Status Canvas_Agent_Service_Poll(
    Canvas_Agent_Service *p_service,
    char *p_response,
    size_t response_capacity,
    char *p_error,
    size_t error_capacity)
{
    Canvas_Agent_Copilot *p_copilot =
        (Canvas_Agent_Copilot *)p_service->p_native;
    if (!p_copilot) return CANVAS_AGENT_ERROR;
    pthread_mutex_lock(&p_copilot->mutex);
    Canvas_Agent_Status status = p_service->status;
    if (status == CANVAS_AGENT_READY) {
        snprintf(p_response, response_capacity, "%s", p_service->response);
        p_service->status = CANVAS_AGENT_IDLE;
    } else if (status == CANVAS_AGENT_ERROR) {
        snprintf(p_error, error_capacity, "%s", p_service->error);
        p_service->status = CANVAS_AGENT_IDLE;
    }
    pthread_mutex_unlock(&p_copilot->mutex);
    return status;
}

void Canvas_Agent_Service_Shutdown(Canvas_Agent_Service *p_service)
{
    Canvas_Agent_Copilot *p_copilot =
        (Canvas_Agent_Copilot *)p_service->p_native;
    if (!p_copilot) return;
    Inference_Bridge_Destroy(p_copilot->p_bridge);
    p_copilot->p_bridge = NULL;
    pthread_mutex_destroy(&p_copilot->mutex);
    p_service->p_native = NULL;
}