"use client";
import * as React from "react";
import { Check, Columns3, Loader2, Play, Square, X } from "lucide-react";
import { useApp } from "@/components/app/store";
import { api, streamEvents, ClientApiError } from "@/lib/client/api";
import type { ChatAdoptResponse, PolyModel, PolyProviderErrorShape } from "@/lib/client/types";
import type { ChatSettings } from "./model-config";
import { ModelSelector } from "./model-selector";
import { MessageError } from "./message-error";
import { Markdown } from "@/components/markdown/markdown";
import { ProviderIcon } from "@/components/brand/provider-icon";
import { Button } from "@/components/ui/button";
import { Textarea } from "@/components/ui/input";
import { Tooltip } from "@/components/ui/tooltip";
import { toast } from "@/components/ui/toast";
import { useIsMobile, useSnapCarousel } from "@/lib/client/hooks";
import { PROVIDERS } from "@/lib/client/providers";
import { cn, formatMs, formatTokens, formatUsd } from "@/lib/utils";

/** Wire shape of an Arena response row (subset we need). */
interface ArenaResponseLite {
  id: string;
  modelKey: string;
  content: string;
  status: string;
  usage?: Record<string, number> | null;
  latencyMs?: number | null;
  ttftMs?: number | null;
  costUsd?: number | null;
  error?: { code: string; message: string } | null;
}

type ArenaEvent =
  | { type: "meta"; responseId: string; modelKey: string }
  | { type: "text-delta"; text: string }
  | { type: "reasoning-delta"; text: string }
  | { type: "citation"; citation: { url?: string; title?: string; snippet?: string } }
  | { type: "server-tool"; name: string; status: "started" | "completed" }
  | { type: "error"; error: PolyProviderErrorShape }
  | { type: "done"; response: ArenaResponseLite; status: "complete" | "stopped" | "error" };

type Status = "waiting" | "thinking" | "streaming" | "done" | "error" | "stopped";

interface Col {
  modelKey: string;
  responseId: string | null;
  status: Status;
  text: string;
  reasoning: string;
  response: ArenaResponseLite | null;
  error: PolyProviderErrorShape | { code: string; message: string } | null;
  errorAt?: number;
  startedAt: number;
  firstTokenAt?: number;
}

export interface CompareInlineProps {
  prompt: string;
  /** Editable prompt (new chat / empty state). */
  onPromptChange?: (p: string) => void;
  systemPrompt?: string | null;
  settings?: ChatSettings;
  attachmentIds?: string[];
  initialModelKeys: string[];
  conversationId?: string | null;
  projectId?: string | null;
  onClose: () => void;
  /** Called after "Continue with <model>" succeeded. */
  onAdopted: (res: ChatAdoptResponse, modelKey: string) => void;
}

const MAX = 4;
const MIN = 2;

function cleanSettings(s: ChatSettings | undefined): ChatSettings | undefined {
  if (!s) return undefined;
  const { tools: _tools, ...rest } = s;
  void _tools;
  const out = Object.fromEntries(Object.entries(rest).filter(([, v]) => v !== undefined && v !== null && !(Array.isArray(v) && v.length === 0))) as ChatSettings;
  return Object.keys(out).length ? out : undefined;
}

function emptyCol(modelKey: string): Col {
  return { modelKey, responseId: null, status: "waiting", text: "", reasoning: "", response: null, error: null, startedAt: Date.now() };
}

/**
 * Inline "Compare with…" — 2–4 models answer the same prompt inside the chat. Runs through the Arena API
 * (so the comparison is stored as an Arena session) and lets the user continue the conversation with any
 * answer via POST /api/chat/adopt.
 */
export function CompareInline({ prompt, onPromptChange, systemPrompt, settings, attachmentIds, initialModelKeys, conversationId, projectId, onClose, onAdopted }: CompareInlineProps) {
  const { modelsByKey, connectedProviders, preferences } = useApp();
  const isMobile = useIsMobile();
  const [selected, setSelected] = React.useState<string[]>(() => initialModelKeys.filter((k) => modelsByKey.has(k)).slice(0, MAX));
  const [columns, setColumns] = React.useState<Col[]>([]);
  const [running, setRunning] = React.useState(false);
  const [adopting, setAdopting] = React.useState<string | null>(null);
  const mapRef = React.useRef(new Map<string, Col>());
  const timer = React.useRef<ReturnType<typeof setTimeout> | null>(null);
  const abortRef = React.useRef<AbortController | null>(null);
  const { ref: rowRef, index, scrollTo } = useSnapCarousel<HTMLDivElement>(columns.length);

  const flush = React.useCallback(() => {
    timer.current = null;
    setColumns([...mapRef.current.values()].map((c) => ({ ...c })));
  }, []);
  const schedule = React.useCallback(() => {
    if (!timer.current) timer.current = setTimeout(flush, 40);
  }, [flush]);
  React.useEffect(
    () => () => {
      if (timer.current) clearTimeout(timer.current);
      abortRef.current?.abort();
    },
    [],
  );

  const selectedSet = React.useMemo(() => new Set(selected), [selected]);
  const toggle = (key: string) =>
    setSelected((s) => {
      if (s.includes(key)) return s.filter((k) => k !== key);
      if (s.length >= MAX) {
        toast.warning(`Up to ${MAX} models`, "Remove one before adding another.");
        return s;
      }
      return [...s, key];
    });
  const disconnected = selected.map((k) => modelsByKey.get(k)).filter((m): m is PolyModel => Boolean(m) && !connectedProviders.has(m!.provider));
  const canRun = !running && selected.length >= MIN && prompt.trim().length > 0 && disconnected.length === 0;
  const reason = running ? null : selected.length < MIN ? `Pick at least ${MIN} models` : !prompt.trim() ? "Write a prompt first" : disconnected.length ? `${disconnected.map((m) => PROVIDERS[m.provider].shortName).join(", ")} not connected` : null;

  const streamOne = async (sessionId: string, modelKey: string, signal: AbortSignal) => {
    const col = () => mapRef.current.get(modelKey);
    try {
      await streamEvents(
        "/api/arena/stream",
        { sessionId, modelKey },
        (raw) => {
          const c = col();
          if (!c) return;
          const ev = raw as unknown as ArenaEvent;
          switch (ev.type) {
            case "meta":
              c.responseId = ev.responseId;
              break;
            case "text-delta":
              if (!c.firstTokenAt) c.firstTokenAt = Date.now();
              c.text += ev.text;
              c.status = "streaming";
              break;
            case "reasoning-delta":
              if (!c.firstTokenAt) c.firstTokenAt = Date.now();
              c.reasoning += ev.text;
              if (c.status === "waiting") c.status = "thinking";
              break;
            case "error":
              c.error = ev.error;
              c.errorAt = Date.now();
              break;
            case "done":
              c.response = ev.response;
              c.responseId = ev.response.id;
              c.status = ev.status === "complete" ? "done" : ev.status === "stopped" ? "stopped" : "error";
              if (ev.response.error && !c.error) {
                c.error = ev.response.error;
                c.errorAt = Date.now();
              }
              break;
            default:
              break;
          }
          schedule();
        },
        signal,
      );
      const c = col();
      if (c && (c.status === "waiting" || c.status === "thinking" || c.status === "streaming")) c.status = c.error ? "error" : "stopped";
    } catch (e) {
      const c = col();
      if (!c) return;
      if ((e as Error).name === "AbortError") c.status = "stopped";
      else {
        const err = e as ClientApiError;
        c.status = "error";
        c.error = { code: err.code ?? "HTTP_ERROR", message: err.message };
        c.errorAt = Date.now();
      }
    }
    flush();
  };

  const run = async () => {
    if (!canRun) return;
    setRunning(true);
    const keys = [...selected];
    mapRef.current = new Map(keys.map((k) => [k, emptyCol(k)]));
    flush();
    const controller = new AbortController();
    abortRef.current = controller;
    try {
      const res = await api<{ session: { id: string } }>("/api/arena", {
        method: "POST",
        json: { prompt: prompt.trim(), systemPrompt: systemPrompt?.trim() || undefined, modelKeys: keys, settings: cleanSettings(settings), attachmentIds: attachmentIds?.length ? attachmentIds.slice(0, 6) : undefined },
      });
      await Promise.allSettled(keys.map((k) => streamOne(res.session.id, k, controller.signal)));
    } catch (e) {
      const err = e as ClientApiError;
      toast.error("Could not start the comparison", err.message);
      mapRef.current = new Map();
      flush();
    } finally {
      setRunning(false);
      abortRef.current = null;
    }
  };

  const stop = () => abortRef.current?.abort();

  const adopt = async (c: Col) => {
    if (!c.responseId && !c.text) return;
    setAdopting(c.modelKey);
    try {
      const res = await api<ChatAdoptResponse>("/api/chat/adopt", {
        method: "POST",
        json: { conversationId: conversationId ?? undefined, modelKey: c.modelKey, arenaResponseId: c.responseId ?? undefined, content: c.responseId ? undefined : c.text, userText: conversationId ? undefined : prompt.trim(), systemPrompt: conversationId ? undefined : systemPrompt ?? null, settings: conversationId ? undefined : cleanSettings(settings), projectId: conversationId ? undefined : projectId ?? null },
      });
      onAdopted(res, c.modelKey);
    } catch (e) {
      toast.error("Could not continue with this answer", (e as Error).message);
    } finally {
      setAdopting(null);
    }
  };

  const allSettled = columns.length > 0 && columns.every((c) => c.status === "done" || c.status === "error" || c.status === "stopped");

  return (
    <section className="rounded-2xl border border-border bg-bg-elevated shadow-sm" aria-label="Compare models" data-no-edge-swipe>
      {/* Header */}
      <header className="flex flex-wrap items-center gap-2 border-b border-border px-3 py-2">
        <span className="flex size-7 items-center justify-center rounded-md bg-accent-soft text-accent">
          <Columns3 className="size-4" />
        </span>
        <h3 className="text-[14px] font-semibold tracking-tight">Compare with…</h3>
        <span className="hidden text-[12px] text-fg-muted sm:inline">{selected.length}/{MAX} models</span>
        <div className="ml-auto flex items-center gap-1.5">
          {running ? (
            <Button size="sm" variant="secondary" onClick={stop}>
              <Square className="size-3.5 fill-current" /> Stop
            </Button>
          ) : (
            <Tooltip content={reason ?? "Run the comparison"}>
              <span className="inline-flex">
                <Button size="sm" variant="accent" disabled={!canRun} onClick={run}>
                  <Play className="size-3.5 fill-current" /> {columns.length ? "Run again" : "Run"}
                </Button>
              </span>
            </Tooltip>
          )}
          <Button size="icon-sm" variant="ghost" onClick={onClose} aria-label="Close comparison" disabled={running}>
            <X />
          </Button>
        </div>
      </header>

      {/* Setup */}
      <div className="space-y-2 px-3 py-2.5">
        <div className="flex flex-wrap items-center gap-1.5">
          {selected.map((k) => {
            const m = modelsByKey.get(k);
            return (
              <span key={k} className="inline-flex h-8 items-center gap-1.5 rounded-full border border-border bg-bg-subtle pl-2 pr-1 text-[12.5px]">
                <ProviderIcon provider={m?.provider ?? k.split("/")[0]} size={13} />
                <span className="max-w-[140px] truncate">{m?.displayName ?? k}</span>
                <button type="button" onClick={() => toggle(k)} disabled={running} className="tap rounded-full p-0.5 text-fg-subtle hover:text-danger" aria-label={`Remove ${m?.displayName ?? k}`}>
                  <X className="size-3" />
                </button>
              </span>
            );
          })}
          {!running ? <ModelSelector value={null} multiple selected={selectedSet} onToggle={toggle} size="sm" buttonLabel={selected.length < MAX ? "+ Add model" : "Models"} className="h-8 rounded-full" /> : null}
        </div>
        {onPromptChange ? (
          <Textarea value={prompt} onChange={(e) => onPromptChange(e.target.value)} placeholder="One prompt for every model…" className="min-h-[64px] text-[14px]" disabled={running} aria-label="Comparison prompt" />
        ) : (
          <p className="line-clamp-3 rounded-lg bg-bg-subtle px-3 py-2 text-[13px] text-fg-muted" title={prompt}>
            {prompt}
          </p>
        )}
        {reason && !columns.length ? <p className="text-[12px] text-fg-subtle">{reason}.</p> : null}
      </div>

      {/* Results */}
      {columns.length ? (
        isMobile ? (
          <div className="border-t border-border">
            <div className="sticky top-0 z-10 flex gap-1 overflow-x-auto bg-bg-elevated/95 px-2 py-2 backdrop-blur scrollbar-none" role="tablist">
              {columns.map((c, i) => {
                const m = modelsByKey.get(c.modelKey);
                return (
                  <button key={c.modelKey} role="tab" aria-selected={index === i} onClick={() => scrollTo(i)} className={cn("inline-flex h-9 shrink-0 items-center gap-1.5 rounded-full border px-3 text-[13px] font-medium", index === i ? "border-fg bg-fg text-bg" : "border-border text-fg-muted")}>
                    <StatusDot status={c.status} />
                    <span className="max-w-[140px] truncate">{m?.displayName ?? c.modelKey}</span>
                  </button>
                );
              })}
            </div>
            <div ref={rowRef} className="snap-row">
              {columns.map((c) => (
                <div key={c.modelKey} className="px-3 pb-3">
                  <Panel col={c} model={modelsByKey.get(c.modelKey)} wrapCode={preferences.codeWrap} showCosts={preferences.showCosts} onAdopt={allSettled ? () => adopt(c) : undefined} adopting={adopting === c.modelKey} />
                </div>
              ))}
            </div>
          </div>
        ) : (
          <div className={cn("grid gap-3 border-t border-border p-3", columns.length === 2 ? "md:grid-cols-2" : columns.length === 3 ? "md:grid-cols-3" : "md:grid-cols-2 xl:grid-cols-4")}>
            {columns.map((c) => (
              <Panel key={c.modelKey} col={c} model={modelsByKey.get(c.modelKey)} wrapCode={preferences.codeWrap} showCosts={preferences.showCosts} onAdopt={allSettled ? () => adopt(c) : undefined} adopting={adopting === c.modelKey} />
            ))}
          </div>
        )
      ) : null}
      {allSettled ? <p className="px-3 pb-3 text-[12px] text-fg-subtle">Saved to the Arena history. Pick an answer to continue the conversation with that model.</p> : null}
    </section>
  );
}

const STATUS: Record<Status, { label: string; dot: string; pulse: boolean }> = {
  waiting: { label: "Waiting", dot: "bg-fg-subtle", pulse: true },
  thinking: { label: "Thinking", dot: "bg-accent", pulse: true },
  streaming: { label: "Streaming", dot: "bg-info", pulse: true },
  done: { label: "Done", dot: "bg-success", pulse: false },
  error: { label: "Failed", dot: "bg-danger", pulse: false },
  stopped: { label: "Stopped", dot: "bg-warning", pulse: false },
};

function StatusDot({ status }: { status: Status }) {
  const s = STATUS[status];
  return <span className={cn("inline-block size-2 shrink-0 rounded-full", s.dot, s.pulse && "animate-pulse-soft")} aria-hidden />;
}

function Panel({ col: c, model, wrapCode, showCosts, onAdopt, adopting }: { col: Col; model?: PolyModel; wrapCode?: boolean; showCosts?: boolean; onAdopt?: () => void; adopting?: boolean }) {
  const name = model?.displayName ?? c.modelKey.split("/").slice(1).join("/");
  const live = c.status === "waiting" || c.status === "thinking" || c.status === "streaming";
  const u = (c.response?.usage ?? null) as { inputTokens?: number; outputTokens?: number } | null;
  const gen = c.response?.latencyMs ? Math.max(1, c.response.latencyMs - (c.response.ttftMs ?? 0)) : null;
  const tps = u?.outputTokens && gen ? Math.round((u.outputTokens / gen) * 1000) : null;
  return (
    <article className="flex min-w-0 flex-col overflow-hidden rounded-xl border border-border bg-bg" aria-label={`${name} response`} aria-busy={live}>
      <header className="flex items-center gap-2 border-b border-border px-3 py-2">
        <ProviderIcon provider={model?.provider ?? c.modelKey.split("/")[0]} size={14} />
        <span className="min-w-0 flex-1 truncate text-[13px] font-semibold">{name}</span>
        <span className="inline-flex items-center gap-1.5 text-[11px] text-fg-subtle">
          <StatusDot status={c.status} /> {STATUS[c.status].label}
        </span>
      </header>
      <div className="max-h-[420px] min-h-[96px] overflow-y-auto px-3 py-2 text-[14px] scrollbar-thin">
        {c.status === "waiting" && !c.text ? (
          <p className="flex items-center gap-2 text-[13px] text-fg-muted">
            <Loader2 className="size-3.5 animate-spin" /> Waiting for the model…
          </p>
        ) : null}
        {c.status === "thinking" && !c.text ? <p className="text-[13px] text-fg-muted">Thinking…</p> : null}
        {c.text ? <Markdown content={c.text} wrap={wrapCode} streaming={live} className="text-[14px]" /> : null}
        {c.error && !live ? (
          <div className="mt-2">
            <MessageError error={c.error} provider={model?.provider ?? c.modelKey.split("/")[0]} since={c.errorAt} compact />
          </div>
        ) : null}
      </div>
      <footer className="flex flex-wrap items-center gap-x-2 gap-y-1 border-t border-border px-3 py-2 text-[11px] tabular-nums text-fg-subtle">
        {c.response?.ttftMs ? <span title="Time to first token">TTFT {formatMs(c.response.ttftMs)}</span> : null}
        {c.response?.latencyMs ? <span>{formatMs(c.response.latencyMs)}</span> : null}
        {tps ? <span>{tps} tok/s</span> : null}
        {u ? (
          <span>
            {formatTokens(u.inputTokens ?? 0)} in · {formatTokens(u.outputTokens ?? 0)} out
          </span>
        ) : null}
        {showCosts && c.response?.costUsd !== null && c.response?.costUsd !== undefined ? <span>≈ {formatUsd(c.response.costUsd, { precise: c.response.costUsd < 0.01 })}</span> : null}
        <span className="flex-1" />
        {onAdopt && c.status === "done" && c.text ? (
          <Button size="xs" variant="outline" className="bg-bg-elevated" onClick={onAdopt} loading={adopting}>
            <Check /> Continue with {name}
          </Button>
        ) : null}
      </footer>
    </article>
  );
}