import { NextRequest } from "next/server";
import { auth } from "@/lib/auth";
import { prisma } from "@/lib/prisma";
import { pullModelStream, parseNdjson, type OllamaPullChunk } from "@/lib/ollama";

export const runtime = "nodejs";

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

  const body = (await req.json().catch(() => null)) as { name?: string } | null;
  if (!body?.name) return new Response("Missing model name", { status: 400 });

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

  const encoder = new TextEncoder();
  const stream = new ReadableStream({
    async start(controller) {
      const send = (data: unknown) => {
        controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`));
      };
      try {
        const ndjson = await pullModelStream(user.ollamaUrl, body.name!, req.signal);
        for await (const chunk of parseNdjson<OllamaPullChunk>(ndjson)) {
          send(chunk);
          if (chunk.error) break;
          if (chunk.status === "success") break;
        }
        send({ status: "done" });
      } catch (e) {
        send({ error: e instanceof Error ? e.message : "Pull failed" });
      } finally {
        controller.close();
      }
    },
  });

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