import { NextRequest } from "next/server";
import { auth } from "@/lib/auth";
import { prisma } from "@/lib/prisma";
import {
  showModel,
  supportsThinking,
  supportsTools,
  type OllamaChatMessage,
  type OllamaToolCall,
} from "@/lib/ollama";
import {
  getToolDefinitions,
  TOOL_SYSTEM_PROMPT,
  ALL_TOOL_NAMES,
  type ToolContext,
} from "@/lib/tools";
import { runAgentLoop, type PersistenceLayer } from "@/lib/chat-engine";
import { z } from "zod";

export const runtime = "nodejs";
export const maxDuration = 600;

const BodySchema = z.object({
  chatId: z.string(),
  userMessage: z.string().min(1),
});

interface StoredMessage {
  role: string;
  content: string;
  toolCalls: string | null;
  toolName: string | null;
  toolCallId: string | null;
}

function buildHistory(messages: StoredMessage[], systemPrompt: string): OllamaChatMessage[] {
  const out: OllamaChatMessage[] = [];
  if (systemPrompt.trim().length > 0) {
    out.push({ role: "system", content: systemPrompt });
  }
  for (const m of messages) {
    if (m.role === "system") continue;
    if (m.role === "tool") {
      out.push({
        role: "tool",
        content: m.content,
        ...(m.toolCallId ? { tool_call_id: m.toolCallId } : {}),
        ...(m.toolName ? { name: m.toolName } : {}),
      });
      continue;
    }
    const msg: OllamaChatMessage = {
      role: m.role === "assistant" ? "assistant" : "user",
      content: m.content,
    };
    if (m.role === "assistant" && m.toolCalls) {
      try {
        const parsed = JSON.parse(m.toolCalls) as OllamaToolCall[];
        if (Array.isArray(parsed) && parsed.length > 0) msg.tool_calls = parsed;
      } catch {
        /* ignore */
      }
    }
    out.push(msg);
  }
  return out;
}

export async function POST(req: NextRequest) {
  const session = await auth();
  if (!session?.user?.id) return new Response("Unauthorized", { status: 401 });

  const json = await req.json().catch(() => null);
  const parsed = BodySchema.safeParse(json);
  if (!parsed.success) return new Response("Invalid body", { status: 400 });

  const { chatId, userMessage } = parsed.data;

  const chat = await prisma.chat.findUnique({
    where: { id: chatId },
    include: { messages: { orderBy: { createdAt: "asc" } } },
  });
  if (!chat || chat.userId !== session.user.id)
    return new Response("Not found", { status: 404 });

  const user = await prisma.user.findUnique({ where: { id: session.user.id } });
  if (!user) return new Response("User not found", { status: 404 });

  // Persist user message immediately so it's visible if the request is retried/refreshed
  await prisma.message.create({
    data: { chatId, role: "user", content: userMessage },
  });

  let canThink = false;
  let canTools = false;
  try {
    const info = await showModel(user.ollamaUrl, chat.model);
    canThink = supportsThinking(info);
    canTools = supportsTools(chat.model, info);
  } catch {
    /* best effort */
  }

  const enabledToolNames = (user.enabledTools || "")
    .split(",")
    .map((s) => s.trim())
    .filter((s) => ALL_TOOL_NAMES.includes(s));
  const toolsActive = canTools && user.toolsEnabled && enabledToolNames.length > 0;
  const tools = toolsActive ? getToolDefinitions(enabledToolNames) : undefined;
  const useThink = canThink && chat.showThinking;

  const systemPrompt = toolsActive
    ? `${TOOL_SYSTEM_PROMPT}${chat.systemPrompt ? "\n\n" + chat.systemPrompt : ""}`
    : chat.systemPrompt;

  const allMessages: StoredMessage[] = [
    ...chat.messages.map((m) => ({
      role: m.role,
      content: m.content,
      toolCalls: m.toolCalls,
      toolName: m.toolName,
      toolCallId: m.toolCallId,
    })),
    { role: "user", content: userMessage, toolCalls: null, toolName: null, toolCallId: null },
  ];
  const history = buildHistory(allMessages, systemPrompt);

  const toolContext: ToolContext = {
    userId: user.id,
    searxngUrl: user.searxngUrl,
  };

  const persistence: PersistenceLayer = {
    async saveAssistant(d) {
      const m = await prisma.message.create({
        data: {
          chatId,
          role: "assistant",
          content: d.content,
          thinking: d.thinking || null,
          toolCalls: d.toolCalls.length > 0 ? JSON.stringify(d.toolCalls) : null,
          promptTokens: d.promptTokens ?? null,
          tokens: d.tokens ?? null,
          evalDurationMs: d.evalDurationMs,
          totalDurationMs: d.totalDurationMs,
        },
      });
      return { id: m.id };
    },
    async saveToolResult(d) {
      const m = await prisma.message.create({
        data: {
          chatId,
          role: "tool",
          content: d.content,
          toolName: d.toolName,
          toolCallId: d.toolCallId,
          toolMeta: d.toolMeta ? JSON.stringify(d.toolMeta) : null,
          totalDurationMs: d.durationMs,
        },
      });
      return { id: m.id };
    },
    async bumpChat() {
      await prisma.chat.update({ where: { id: chatId }, data: { updatedAt: new Date() } });
    },
  };

  const encoder = new TextEncoder();
  const stream = new ReadableStream({
    async start(controller) {
      let closed = false;
      const send = (evt: { type: string; data?: unknown }) => {
        if (closed) return;
        controller.enqueue(encoder.encode(`data: ${JSON.stringify(evt)}\n\n`));
      };
      try {
        await runAgentLoop({
          ollamaUrl: user.ollamaUrl,
          toolContext,
          model: chat.model,
          contextWindow: chat.contextWindow,
          temperature: chat.temperature,
          keepAlive: chat.keepAlive,
          useThink,
          tools,
          history,
          maxIterations: user.maxToolIterations,
          signal: req.signal,
          send,
          persistence,
        });
      } catch (e) {
        send({ type: "error", data: e instanceof Error ? e.message : "Stream failed" });
      } finally {
        closed = true;
        controller.close();
      }
    },
  });

  return new Response(stream, {
    headers: {
      "content-type": "text/event-stream",
      "cache-control": "no-cache, no-transform",
      connection: "keep-alive",
    },
  });
}