// app/api/chat/route.ts
import { NextRequest } from "next/server";

export const runtime = "edge";

export async function POST(req: NextRequest) {
  const { messages } = await req.json();

  const upstream = await fetch("https://api.siati.ai/v1/chat/completions", {
    method: "POST",
    headers: {
      "Authorization": `Bearer ${process.env.SIATI_API_KEY}`,
      "Content-Type":  "application/json",
    },
    body: JSON.stringify({
      model: "apertus-70b-instruct",
      messages,
      stream: true,
      stream_options: { include_usage: true },
    }),
  });

  if (!upstream.ok || !upstream.body) {
    return new Response("upstream error", { status: 502 });
  }

  // Forward the SSE stream as-is
  return new Response(upstream.body, {
    headers: {
      "Content-Type":  "text/event-stream",
      "Cache-Control": "no-cache",
      "Connection":    "keep-alive",
    },
  });
}

"use client";
import { useState } from "react";

export function Chat() {
  const [answer, setAnswer] = useState("");
  const [loading, setLoading] = useState(false);

  async function send(question: string) {
    setAnswer("");
    setLoading(true);

    const res = await fetch("/api/chat", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        messages: [{ role: "user", content: question }],
      }),
    });

    if (!res.body) { setLoading(false); return; }

    const reader = res.body.getReader();
    const decoder = new TextDecoder();
    let buf = "";

    try {
      while (true) {
        const { value, done } = await reader.read();
        if (done) break;
        buf += decoder.decode(value, { stream: true });

        const lines = buf.split("\n");
        buf = lines.pop() ?? "";

        for (const line of lines) {
          const t = line.trim();
          if (!t.startsWith("data:")) continue;
          const data = t.slice(5).trim();
          if (data === "[DONE]") { setLoading(false); return; }

          try {
            const chunk = JSON.parse(data);
            const delta = chunk.choices?.[0]?.delta?.content;
            if (delta) setAnswer(prev => prev + delta);
          } catch { /* ignore parse errors mid-stream */ }
        }
      }
    } finally {
      setLoading(false);
    }
  }

  return (
    <div>
      <button onClick={() => send("Conta da 1 a 10.")} disabled={loading}>
        Ask
      </button>
      <pre>{answer}</pre>
    </div>
  );
}

const abortRef = useRef<AbortController | null>(null);

useEffect(() => () => abortRef.current?.abort(), []);

async function send(question: string) {
  abortRef.current?.abort();
  abortRef.current = new AbortController();

  const res = await fetch("/api/chat", {
    signal: abortRef.current.signal,
    /* ... */
  });
  // ...
}

let usage: { prompt_tokens: number; completion_tokens: number } | null = null;
// inside the stream loop:
if (chunk.usage) usage = chunk.usage;
// after [DONE]:
console.log("Used", usage);