import { generate, listGenerationOptions, segment, type ComfyGenerateRequest, type ComfySegmentRequest } from "./comfy";

export async function handleComfyApi(request: Request) {
  try {
    const url = new URL(request.url);
    if (url.pathname === "/api/comfy/models" && request.method === "GET") return json(await listGenerationOptions());
    if (url.pathname === "/api/comfy/generate" && request.method === "POST") {
      const body = await request.json() as ComfyGenerateRequest;
      if (request.headers.get("accept")?.includes("application/x-ndjson")) return generationStream(body, request.signal);
      return json(await generate(body, request.signal));
    }
    if (url.pathname === "/api/comfy/segment" && request.method === "POST") return json(await segment(await request.json() as ComfySegmentRequest, request.signal));
    return new Response("Not found", { status: 404 });
  } catch (error) {
    return new Response(error instanceof Error ? error.message : "ComfyUI request failed", { status: 500 });
  }
}

function generationStream(body: ComfyGenerateRequest, signal: AbortSignal) {
  const encoder = new TextEncoder();
  return new Response(new ReadableStream({
    start(controller) {
      const send = (value: unknown) => controller.enqueue(encoder.encode(`${JSON.stringify(value)}\n`));
      void generate(body, signal, (event) => send({ type: "progress", ...event }))
        .then((result) => { send({ type: "result", ...result }); controller.close(); })
        .catch((error) => { send({ type: "error", message: error instanceof Error ? error.message : "Generation failed" }); controller.close(); });
    },
    cancel() {},
  }), { headers: { "content-type": "application/x-ndjson; charset=utf-8", "cache-control": "no-store" } });
}

function json(value: unknown) {
  return new Response(JSON.stringify(value), { headers: { "content-type": "application/json" } });
}