// Author: Simon-Pierre Boucher
// Contact: contact@spboucher.ai
// Project: chat.spboucher.ai

// Generation runner: explicit state machine per generation
// (queued → starting → streaming → completed | cancelled | failed),
// SSE emission with keep-alives, incremental persistence, cancellation registry.

import crypto from "node:crypto";
import { getDb } from "@/lib/db/database";
import { streamGeneration, estimateCost } from "@/lib/openrouter";
import type { ChatStreamEvent } from "@/lib/openrouter";
import { compileContext } from "@/lib/context";
import {
  getConversation,
  insertMessage,
  threadToLeaf,
  touchConversation,
  updateConversation,
  type MessageRow,
} from "@/lib/conversations";
import { getModel, touchModelUsage, type CatalogModel } from "@/lib/catalog";

const KEEPALIVE_MS = 15_000;
const FLUSH_MS = 700;

// Registry of live generations for cancellation; on globalThis to survive dev HMR.
const registry: Map<string, AbortController> = ((globalThis as Record<string, unknown>).__genRegistry ??=
  new Map()) as Map<string, AbortController>;

export function cancelGeneration(generationId: string): boolean {
  const ctrl = registry.get(generationId);
  if (!ctrl) return false;
  ctrl.abort();
  return true;
}

export function generationState(generationId: string) {
  return getDb()
    .prepare("SELECT id, message_id, state, error_code, error_message, finished_at FROM generations WHERE id = ?")
    .get(generationId) as
    | { id: string; message_id: string; state: string; error_code: string | null; error_message: string | null; finished_at: number | null }
    | undefined;
}

function setGenState(id: string, state: string, errorCode?: string, errorMessage?: string): void {
  const finished = ["completed", "cancelled", "failed"].includes(state) ? Date.now() : null;
  getDb()
    .prepare(
      "UPDATE generations SET state = ?, error_code = COALESCE(?, error_code), error_message = COALESCE(?, error_message), finished_at = COALESCE(?, finished_at) WHERE id = ?"
    )
    .run(state, errorCode ?? null, errorMessage ?? null, finished, id);
}

export interface StartGenerationArgs {
  conversationId: string;
  parentMessage: MessageRow | null; // the user message the assistant answers
  model: CatalogModel;
}

export interface StartedGeneration {
  generationId: string;
  assistantMessage: MessageRow;
  stream: ReadableStream<Uint8Array>;
}

function sseFrame(event: object): string {
  return `data: ${JSON.stringify(event)}\n\n`;
}

/**
 * Create the assistant placeholder + generation row, then return an SSE stream
 * that runs the generation, persists deltas incrementally, and finalizes state.
 */
export function startGeneration(args: StartGenerationArgs): StartedGeneration {
  const db = getDb();
  const generationId = crypto.randomUUID();
  const { conversationId, parentMessage, model } = args;

  const assistantMessage = insertMessage({
    conversationId,
    parentId: parentMessage?.id ?? null,
    role: "assistant",
    modelId: model.id,
    modelName: model.name,
    provider: model.provider,
    generationId,
    status: "pending",
  });

  db.prepare(
    "INSERT INTO generations (id, message_id, conversation_id, model_id, state, created_at) VALUES (?, ?, ?, ?, 'queued', ?)"
  ).run(generationId, assistantMessage.id, conversationId, model.id, Date.now());

  updateConversation(conversationId, { currentLeafId: assistantMessage.id });
  touchConversation(conversationId);
  touchModelUsage(model.id);

  const ctrl = new AbortController();
  registry.set(generationId, ctrl);

  const thread = threadToLeaf(conversationId, parentMessage?.id ?? null);
  const compiled = compileContext(thread, model.contextLength);

  const encoder = new TextEncoder();

  const stream = new ReadableStream<Uint8Array>({
    async start(controller) {
      const send = (e: object) => {
        try {
          controller.enqueue(encoder.encode(sseFrame(e)));
        } catch {
          /* client went away; generation continues until aborted or done */
        }
      };

      // keep-alive comments so ngrok and mobile radios don't kill idle streams
      const keepAlive = setInterval(() => {
        try {
          controller.enqueue(encoder.encode(`: ping\n\n`));
        } catch {
          /* ignore */
        }
      }, KEEPALIVE_MS);

      let content = "";
      let reasoning = "";
      let dirty = false;
      const flush = () => {
        if (!dirty) return;
        db.prepare("UPDATE messages SET content = ?, reasoning = ?, status = 'streaming' WHERE id = ?").run(
          content,
          reasoning || null,
          assistantMessage.id
        );
        dirty = false;
      };
      const flusher = setInterval(flush, FLUSH_MS);

      send({
        type: "meta",
        conversationId,
        assistantMessageId: assistantMessage.id,
        generationId,
        model: model.id,
        estimatedPromptTokens: compiled.estimatedPromptTokens,
      });

      setGenState(generationId, "starting");

      let finalState: "completed" | "cancelled" | "failed" = "completed";
      let errorMessage: string | null = null;
      let sawUsage = false;

      try {
        const events = streamGeneration(
          { model: model.id, messages: compiled.messages, signal: ctrl.signal },
          generationId
        );
        for await (const event of events) {
          switch (event.type) {
            case "generation.start":
              setGenState(generationId, "streaming");
              break;
            case "content.delta":
              content += event.text;
              dirty = true;
              break;
            case "reasoning.delta":
              reasoning += event.text;
              dirty = true;
              break;
            case "usage":
              sawUsage = true;
              recordUsage(generationId, model, event);
              break;
            case "generation.error":
              finalState = "failed";
              errorMessage = event.message;
              break;
            case "generation.end":
            case "tool.start":
            case "tool.delta":
              break;
          }
          send(event);
        }
        if (ctrl.signal.aborted) finalState = "cancelled";
      } catch {
        finalState = ctrl.signal.aborted ? "cancelled" : "failed";
        if (finalState === "failed") errorMessage = "The stream failed unexpectedly.";
        send({ type: "generation.error", message: errorMessage ?? "Cancelled", retryable: finalState === "failed" });
      } finally {
        clearInterval(keepAlive);
        clearInterval(flusher);
        registry.delete(generationId);

        flush();
        db.prepare(
          "UPDATE messages SET content = ?, reasoning = ?, status = ?, error_message = ? WHERE id = ?"
        ).run(content, reasoning || null, finalState, errorMessage, assistantMessage.id);
        setGenState(generationId, finalState, undefined, errorMessage ?? undefined);
        touchConversation(conversationId);

        // If the provider never reported usage, store an estimate for cost tracking.
        if (!sawUsage && finalState === "completed") {
          recordUsage(generationId, model, {
            type: "usage",
            promptTokens: compiled.estimatedPromptTokens,
            completionTokens: Math.ceil(content.length / 3.6),
            totalTokens: compiled.estimatedPromptTokens + Math.ceil(content.length / 3.6),
          });
        }

        send({ type: "state.final", state: finalState, messageId: assistantMessage.id });
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      // Client disconnected (tab closed, network drop). Abort upstream — never keep paying.
      ctrl.abort();
    },
  });

  return { generationId, assistantMessage, stream };
}

function recordUsage(
  generationId: string,
  model: CatalogModel,
  usage: Extract<ChatStreamEvent, { type: "usage" }>
): void {
  const estimated = estimateCost(model.pricing, usage.promptTokens, usage.completionTokens);
  getDb()
    .prepare(
      `INSERT INTO generation_usage
        (generation_id, model_id, prompt_tokens, completion_tokens, reasoning_tokens, cached_tokens, total_tokens, estimated_cost_usd, reported_cost_usd, created_at)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
    )
    .run(
      generationId,
      model.id,
      usage.promptTokens ?? null,
      usage.completionTokens ?? null,
      usage.reasoningTokens ?? null,
      usage.cachedTokens ?? null,
      usage.totalTokens ?? null,
      estimated ?? null,
      usage.cost ?? null,
      Date.now()
    );
}

/** Guard used by the chat route: conversation must exist. */
export function assertConversation(id: string) {
  const conv = getConversation(id);
  if (!conv) throw new Error("Conversation not found");
  return conv;
}