import crypto from "node:crypto";
import {
  chatStream,
  parseNdjson,
  type OllamaChatChunk,
  type OllamaChatMessage,
  type OllamaToolCall,
  type OllamaToolDefinition,
} from "@/lib/ollama";
import { runTool, type ToolContext } from "@/lib/tools";

export type Sender = (evt: { type: string; data?: unknown }) => void;

export interface SaveAssistantInput {
  content: string;
  thinking: string;
  toolCalls: OllamaToolCall[];
  promptTokens?: number;
  tokens?: number;
  evalDurationMs: number;
  totalDurationMs: number;
}

export interface SaveToolInput {
  content: string;
  toolName: string;
  toolCallId: string;
  toolMeta: Record<string, unknown> | null;
  durationMs: number;
}

export interface PersistenceLayer {
  saveAssistant(d: SaveAssistantInput): Promise<{ id: string }>;
  saveToolResult(d: SaveToolInput): Promise<{ id: string }>;
  bumpChat(): Promise<void>;
}

export interface RunAgentInput {
  ollamaUrl: string;
  toolContext: ToolContext;
  model: string;
  contextWindow: number;
  temperature: number;
  keepAlive: number;
  useThink: boolean;
  tools?: OllamaToolDefinition[];
  history: OllamaChatMessage[];
  maxIterations: number;
  signal: AbortSignal;
  send: Sender;
  persistence: PersistenceLayer;
}

/**
 * Persistence layer that throws away all writes. Used by ephemeral chats.
 */
export function noopPersistence(): PersistenceLayer {
  const fakeId = () => `mem_${crypto.randomBytes(6).toString("hex")}`;
  return {
    async saveAssistant() {
      return { id: fakeId() };
    },
    async saveToolResult() {
      return { id: fakeId() };
    },
    async bumpChat() {
      /* noop */
    },
  };
}

export async function runAgentLoop(input: RunAgentInput): Promise<void> {
  const requestStart = Date.now();
  let finalContent = "";
  let finalThinking = "";
  let finalPromptTokens: number | undefined;
  let finalEvalTokens: number | undefined;
  let finalEvalMs: number | undefined;
  let finalTotalMs: number | undefined;

  const { history } = input;
  const maxIter = Math.max(1, Math.min(25, input.maxIterations));

  for (let iter = 0; iter < maxIter; iter++) {
    let iterContent = "";
    let iterThinking = "";
    const iterToolCalls: OllamaToolCall[] = [];
    let iterFinalChunk: OllamaChatChunk | undefined;
    let iterFirstTokenAt: number | undefined;
    let iterContentChunks = 0;
    const iterStart = Date.now();

    const ndjson = await chatStream(
      input.ollamaUrl,
      {
        model: input.model,
        messages: history,
        ...(input.useThink ? { think: true } : {}),
        keep_alive: input.keepAlive,
        ...(input.tools ? { tools: input.tools } : {}),
        options: {
          temperature: input.temperature,
          num_ctx: input.contextWindow,
        },
      },
      input.signal,
    );

    for await (const chunk of parseNdjson<OllamaChatChunk>(ndjson)) {
      if (chunk.message?.thinking) {
        iterThinking += chunk.message.thinking;
        input.send({ type: "thinking", data: chunk.message.thinking });
      }
      if (chunk.message?.content) {
        if (iterFirstTokenAt === undefined) iterFirstTokenAt = Date.now();
        iterContentChunks += 1;
        iterContent += chunk.message.content;
        input.send({ type: "content", data: chunk.message.content });
      }
      if (chunk.message?.tool_calls) {
        for (const tc of chunk.message.tool_calls) iterToolCalls.push(tc);
      }
      if (chunk.done) {
        iterFinalChunk = chunk;
        break;
      }
    }

    // Strip <think>...</think> tags and merge into thinking
    const thinkMatches = [...iterContent.matchAll(/<think>([\s\S]*?)<\/think>/g)];
    if (thinkMatches.length > 0) {
      for (const m of thinkMatches) {
        iterThinking += (iterThinking ? "\n\n" : "") + (m[1] ?? "").trim();
      }
      iterContent = iterContent.replace(/<think>[\s\S]*?<\/think>/g, "").trim();
    }

    const endedAt = Date.now();
    const wallClockMs = endedAt - iterStart;
    const generationMs =
      iterFirstTokenAt !== undefined ? endedAt - iterFirstTokenAt : wallClockMs;
    const evalMs = iterFinalChunk?.eval_duration
      ? Math.round(iterFinalChunk.eval_duration / 1_000_000)
      : generationMs;
    const totalMs = iterFinalChunk?.total_duration
      ? Math.round(iterFinalChunk.total_duration / 1_000_000)
      : wallClockMs;
    let evalTokens = iterFinalChunk?.eval_count;
    if (evalTokens === undefined && iterContentChunks > 0) evalTokens = iterContentChunks;
    const promptTokens = iterFinalChunk?.prompt_eval_count;

    const assistantSaved = await input.persistence.saveAssistant({
      content: iterContent,
      thinking: iterThinking,
      toolCalls: iterToolCalls,
      promptTokens,
      tokens: evalTokens,
      evalDurationMs: evalMs,
      totalDurationMs: totalMs,
    });

    const assistantOllamaMsg: OllamaChatMessage = {
      role: "assistant",
      content: iterContent,
    };
    if (iterToolCalls.length > 0) assistantOllamaMsg.tool_calls = iterToolCalls;
    history.push(assistantOllamaMsg);

    if (iterToolCalls.length === 0) {
      finalContent = iterContent;
      finalThinking = iterThinking;
      finalPromptTokens = promptTokens;
      finalEvalTokens = evalTokens;
      finalEvalMs = evalMs;
      finalTotalMs = Date.now() - requestStart;
      break;
    }

    for (const tc of iterToolCalls) {
      const callId = tc.id || `call_${crypto.randomBytes(6).toString("hex")}`;
      const name = tc.function?.name ?? "";
      const args = tc.function?.arguments ?? {};

      input.send({
        type: "tool_call_start",
        data: { id: callId, messageId: assistantSaved.id, name, arguments: args },
      });

      const toolStart = Date.now();
      const result = await runTool(name, args, input.toolContext);
      const toolDuration = Date.now() - toolStart;

      const toolSaved = await input.persistence.saveToolResult({
        content: result.content,
        toolName: name,
        toolCallId: callId,
        toolMeta: result.meta ?? null,
        durationMs: toolDuration,
      });

      input.send({
        type: "tool_call_result",
        data: {
          id: callId,
          messageId: toolSaved.id,
          name,
          result: result.content,
          error: result.error ?? null,
          meta: result.meta ?? null,
          durationMs: toolDuration,
        },
      });

      history.push({
        role: "tool",
        content: result.content,
        tool_call_id: callId,
        name,
      });
    }

    if (iter === maxIter - 1) {
      input.send({
        type: "warning",
        data: `Limite de ${maxIter} itérations atteinte. Réponse partielle.`,
      });
    }
  }

  await input.persistence.bumpChat();

  input.send({
    type: "done",
    data: {
      promptTokens: finalPromptTokens,
      evalTokens: finalEvalTokens,
      evalDurationMs: finalEvalMs,
      totalDurationMs: finalTotalMs,
      thinking: finalThinking,
      content: finalContent,
    },
  });
}