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, noopPersistence } from "@/lib/chat-engine";
import { z } from "zod";

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

const PriorMessage = z.object({
  role: z.enum(["user", "assistant", "tool"]),
  content: z.string(),
  toolCalls: z
    .array(
      z.object({
        id: z.string().optional(),
        function: z.object({
          name: z.string(),
          arguments: z.union([z.record(z.unknown()), z.string()]).optional(),
        }),
      }),
    )
    .optional()
    .nullable(),
  toolName: z.string().optional().nullable(),
  toolCallId: z.string().optional().nullable(),
});

const BodySchema = z.object({
  userMessage: z.string().min(1),
  model: z.string().min(1),
  contextWindow: z.number().int().positive().default(8192),
  systemPrompt: z.string().default(""),
  temperature: z.number().min(0).max(2).default(0.7),
  showThinking: z.boolean().default(true),
  keepAlive: z.number().int().min(-1).default(300),
  toolsEnabled: z.boolean().default(true),
  history: z.array(PriorMessage).default([]),
});

function buildHistory(
  history: z.infer<typeof BodySchema>["history"],
  systemPrompt: string,
  userMessage: string,
): OllamaChatMessage[] {
  const out: OllamaChatMessage[] = [];
  if (systemPrompt.trim().length > 0) {
    out.push({ role: "system", content: systemPrompt });
  }
  for (const m of history) {
    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, content: m.content };
    if (m.role === "assistant" && m.toolCalls && m.toolCalls.length > 0) {
      msg.tool_calls = m.toolCalls as OllamaToolCall[];
    }
    out.push(msg);
  }
  out.push({ role: "user", content: userMessage });
  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(
      JSON.stringify({ error: parsed.error.issues[0]?.message ?? "Invalid body" }),
      { status: 400, headers: { "content-type": "application/json" } },
    );
  }

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

  const body = parsed.data;

  let canThink = false;
  let canTools = false;
  try {
    const info = await showModel(user.ollamaUrl, body.model);
    canThink = supportsThinking(info);
    canTools = supportsTools(body.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 && body.toolsEnabled && enabledToolNames.length > 0;
  const tools = toolsActive ? getToolDefinitions(enabledToolNames) : undefined;
  const useThink = canThink && body.showThinking;

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

  const history = buildHistory(body.history, systemPrompt, body.userMessage);

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

  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: body.model,
          contextWindow: body.contextWindow,
          temperature: body.temperature,
          keepAlive: body.keepAlive,
          useThink,
          tools,
          history,
          maxIterations: user.maxToolIterations,
          signal: req.signal,
          send,
          persistence: noopPersistence(),
        });
      } 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",
    },
  });
}