"use client";

import { useSearchParams } from "next/navigation";
import { Suspense, useEffect, useMemo, useRef, useState } from "react";
import { api, ApiError, streamChat } from "@/lib/api";
import { useLive } from "@/lib/events";
import { fmtMs } from "@/lib/format";
import type { Model } from "@/lib/types";
import { Code, Modal, Pill, StatusPill } from "@/components/ui";

interface Msg { role: "user" | "assistant" | "system"; content: string; reasoning?: string; stats?: { tps?: number; ttft?: number; tokens?: number; prompt?: number; total?: number } }

function Playground() {
  const params = useSearchParams();
  const live = useLive();
  const [models, setModels] = useState<Model[]>([]);
  const [model, setModel] = useState(params.get("model") || "");
  const [system, setSystem] = useState("");
  const [temp, setTemp] = useState(0.7);
  const [topP, setTopP] = useState(0.95);
  const [maxTokens, setMaxTokens] = useState(1024);
  const [thinking, setThinking] = useState(true);
  const [msgs, setMsgs] = useState<Msg[]>([]);
  const [input, setInput] = useState("");
  const [busy, setBusy] = useState(false);
  const [phase, setPhase] = useState<string>("");
  const [showRaw, setShowRaw] = useState(false);
  const [showSettings, setShowSettings] = useState(false);
  const abortRef = useRef<AbortController | null>(null);
  const endRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    api.get<{ models: Model[] }>("/api/models").then((r) => {
      const ok = r.models.filter((m) => m.installed && m.enabled && !m.embedding && !m.reranker);
      setModels(ok);
      if (!model && ok.length) {
        const cur = ok.find((m) => m.status === "ready") || ok.find((m) => m.favorite) || ok[0];
        setModel(cur.id);
      }
    }).catch(() => {});
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [live.version]);

  useEffect(() => { endRef.current?.scrollIntoView({ block: "end" }); }, [msgs, phase]);

  const current = models.find((m) => m.id === model);
  const status = live.manager?.loaded.find((w) => w.model_id === model)?.status || live.manager?.progress[model]?.status || current?.status || "unloaded";

  const body = useMemo(() => {
    const messages: { role: string; content: string }[] = [];
    if (system.trim()) messages.push({ role: "system", content: system });
    for (const m of msgs) if (m.role !== "system") messages.push({ role: m.role, content: m.content });
    const b: Record<string, unknown> = { model, messages, temperature: temp, top_p: topP, max_tokens: maxTokens, stream: true };
    if (current?.thinking) b.chat_template_kwargs = { enable_thinking: thinking };
    return b;
  }, [model, system, msgs, temp, topP, maxTokens, thinking, current]);

  const send = async (text?: string, regenerate = false) => {
    const userText = text ?? input.trim();
    if (!model || busy) return;
    let history = msgs;
    if (regenerate) {
      history = msgs.slice(0, msgs.length - 1);
    } else {
      if (!userText) return;
      history = [...msgs, { role: "user", content: userText }];
      setInput("");
    }
    const assistant: Msg = { role: "assistant", content: "", reasoning: "" };
    setMsgs([...history, assistant]);
    setBusy(true);
    setPhase(status === "ready" ? "Generating…" : "Loading model…");
    const ac = new AbortController();
    abortRef.current = ac;
    const t0 = performance.now();
    let first = 0;
    try {
      const messages: { role: string; content: string }[] = [];
      if (system.trim()) messages.push({ role: "system", content: system });
      for (const m of history) messages.push({ role: m.role, content: m.content });
      const req: Record<string, unknown> = { model, messages, temperature: temp, top_p: topP, max_tokens: maxTokens };
      if (current?.thinking) req.chat_template_kwargs = { enable_thinking: thinking };
      for await (const chunk of streamChat(req, ac.signal)) {
        if ((chunk as { error?: { message: string } }).error) throw new Error((chunk as { error: { message: string } }).error.message);
        const choices = (chunk.choices as { delta?: { content?: string; reasoning_content?: string } }[]) || [];
        const d = choices[0]?.delta;
        if (d?.content || d?.reasoning_content) {
          if (!first) { first = performance.now(); setPhase("Generating…"); }
          assistant.content += d.content || "";
          assistant.reasoning = (assistant.reasoning || "") + (d.reasoning_content || "");
          setMsgs([...history, { ...assistant }]);
        }
        const usage = chunk.usage as { completion_tokens?: number; prompt_tokens?: number } | undefined;
        const timings = chunk.timings as { generation_tps?: number; ttft_ms?: number } | undefined;
        if (usage || timings) {
          assistant.stats = {
            tokens: usage?.completion_tokens, prompt: usage?.prompt_tokens,
            tps: timings?.generation_tps, ttft: timings?.ttft_ms ?? (first ? first - t0 : undefined), total: performance.now() - t0,
          };
          setMsgs([...history, { ...assistant }]);
        }
      }
    } catch (e) {
      if ((e as Error).name !== "AbortError") {
        assistant.content += `\n\n⚠️ ${e instanceof ApiError ? e.message : (e as Error).message}`;
        setMsgs([...history, { ...assistant }]);
      }
    } finally {
      if (!assistant.stats) assistant.stats = { total: performance.now() - t0, ttft: first ? first - t0 : undefined };
      setMsgs([...history, { ...assistant }]);
      setBusy(false);
      setPhase("");
      abortRef.current = null;
    }
  };

  return (
    <div className="flex flex-col h-[calc(100vh-56px)] md:h-[calc(100vh-64px)] -mb-20 md:mb-0">
      <div className="flex flex-wrap items-center gap-2 pb-3 border-b border-border">
        <select className="input w-auto max-w-[280px]" value={model} onChange={(e) => setModel(e.target.value)}>
          {!models.length && <option value="">No models installed</option>}
          {models.map((m) => <option key={m.id} value={m.id}>{m.favorite ? "★ " : ""}{m.name}{m.status === "ready" ? " · loaded" : ""}</option>)}
        </select>
        <StatusPill status={status} />
        {current && <span className="text-xs text-ink-3 hidden sm:inline">{current.runtime === "mlx" ? "MLX" : "llama.cpp"} · {current.quantization} · ctx {Math.round((current.recommended_context || 0) / 1024)}K</span>}
        <div className="ml-auto flex gap-2">
          <button className="btn btn-sm" onClick={() => setShowSettings(true)}>Parameters</button>
          <button className="btn btn-sm" onClick={() => setShowRaw(true)}>Raw request</button>
          <button className="btn btn-sm btn-ghost" onClick={() => setMsgs([])} disabled={busy}>Clear</button>
        </div>
      </div>

      <div className="flex-1 overflow-auto py-4 flex flex-col gap-4">
        {!msgs.length && (
          <div className="m-auto text-center text-ink-3 text-sm max-w-md">
            <div className="text-2xl mb-2">λ</div>
            Send a message. If <span className="text-ink">{current?.name || "the model"}</span> is not loaded, it is loaded from SSD first — you will see the status change above.
          </div>
        )}
        {msgs.map((m, i) => (
          <div key={i} className={`flex ${m.role === "user" ? "justify-end" : "justify-start"}`}>
            <div className={`max-w-[85%] md:max-w-[75%] rounded-2xl px-4 py-2.5 text-[14px] leading-relaxed whitespace-pre-wrap ${m.role === "user" ? "bg-accent text-white rounded-br-md" : "card rounded-bl-md"}`}>
              {m.reasoning && (
                <details className="mb-2 text-xs text-ink-3">
                  <summary className="cursor-pointer select-none">Reasoning ({m.reasoning.length} chars)</summary>
                  <div className="mt-1 whitespace-pre-wrap border-l-2 border-border pl-2 max-h-60 overflow-auto">{m.reasoning}</div>
                </details>
              )}
              {m.content || (busy && i === msgs.length - 1 ? <span className="text-ink-3 pulse">{phase || "…"}</span> : "")}
              {m.role === "assistant" && m.stats && !(busy && i === msgs.length - 1) && (
                <div className="mt-2 pt-2 border-t border-border flex flex-wrap gap-x-3 gap-y-1 text-[11px] text-ink-3 num">
                  {m.stats.tokens != null && <span>{m.stats.tokens} tokens</span>}
                  {m.stats.tps && <span>{m.stats.tps.toFixed(1)} tok/s</span>}
                  {m.stats.ttft != null && <span>TTFT {fmtMs(m.stats.ttft)}</span>}
                  {m.stats.total != null && <span>{fmtMs(m.stats.total)} total</span>}
                  <button className="hover:text-ink" onClick={() => navigator.clipboard.writeText(m.content)}>copy</button>
                  {i === msgs.length - 1 && <button className="hover:text-ink" onClick={() => send(undefined, true)}>regenerate</button>}
                </div>
              )}
            </div>
          </div>
        ))}
        <div ref={endRef} />
      </div>

      <form className="pt-3 border-t border-border flex gap-2 items-end" onSubmit={(e) => { e.preventDefault(); send(); }}>
        <textarea className="input flex-1 min-h-[44px] max-h-40" rows={1} placeholder={`Message ${current?.name || "model"}… (⌘/Ctrl+Enter to send)`} value={input}
          onChange={(e) => setInput(e.target.value)}
          onKeyDown={(e) => { if ((e.metaKey || e.ctrlKey) && e.key === "Enter") { e.preventDefault(); send(); } }} />
        {busy ? (
          <button type="button" className="btn btn-danger h-[44px]" onClick={() => abortRef.current?.abort()}>Stop</button>
        ) : (
          <button className="btn btn-primary h-[44px]" disabled={!input.trim() || !model}>Send</button>
        )}
      </form>

      <Modal open={showSettings} onClose={() => setShowSettings(false)} title="Parameters">
        <div className="flex flex-col gap-4 text-sm">
          <label className="flex flex-col gap-1"><span className="label">System prompt</span><textarea className="input" rows={3} value={system} onChange={(e) => setSystem(e.target.value)} placeholder="You are a helpful assistant." /></label>
          <label className="flex flex-col gap-1"><span className="label flex justify-between">Temperature <span className="num">{temp.toFixed(2)}</span></span><input type="range" min={0} max={2} step={0.05} value={temp} onChange={(e) => setTemp(Number(e.target.value))} /></label>
          <label className="flex flex-col gap-1"><span className="label flex justify-between">Top-p <span className="num">{topP.toFixed(2)}</span></span><input type="range" min={0.05} max={1} step={0.05} value={topP} onChange={(e) => setTopP(Number(e.target.value))} /></label>
          <label className="flex flex-col gap-1"><span className="label">Max tokens</span><input className="input" type="number" min={1} max={32768} value={maxTokens} onChange={(e) => setMaxTokens(Number(e.target.value))} /></label>
          {current?.thinking && <label className="flex items-center gap-2"><input type="checkbox" checked={thinking} onChange={(e) => setThinking(e.target.checked)} /> Enable thinking (reasoning models)</label>}
          <div className="text-xs text-ink-3">Context window of the loaded worker: {current?.recommended_context ? `${Math.round(current.recommended_context / 1024)}K tokens` : "—"} <Pill>{current?.runtime}</Pill></div>
        </div>
      </Modal>
      <Modal open={showRaw} onClose={() => setShowRaw(false)} title="Raw API request" width={680}>
        <Code>{`curl ${typeof location !== "undefined" ? location.origin : ""}/v1/chat/completions \\
  -H "Authorization: Bearer llm_live_..." \\
  -H "Content-Type: application/json" \\
  -d '${JSON.stringify(body, null, 2)}'`}</Code>
      </Modal>
    </div>
  );
}

export default function PlaygroundPage() {
  return <Suspense><Playground /></Suspense>;
}