"use client";
import { useCallback, useEffect, useSyncExternalStore } from "react";
import { api, readSSE } from "./fetcher";
import type { MessageDTO, ConversationDTO, ChatSendOptions, Part, AttachmentDTO, ArtifactDTO } from "./types";
import { toast } from "@/components/ui";
import { mutate } from "swr";

export interface ChatState {
  conversationId: string | null;
  conversation: ConversationDTO | null;
  messages: MessageDTO[];
  allMessages: MessageDTO[];
  siblings: { parent_id: string | null; role: string; n: number }[];
  attachments: AttachmentDTO[];
  artifacts: ArtifactDTO[];
  loading: boolean;
  streaming: boolean;
  error: string | null;
  warning: string | null;
  abort: AbortController | null;
}

const empty = (id: string | null): ChatState => ({ conversationId: id, conversation: null, messages: [], allMessages: [], siblings: [], attachments: [], artifacts: [], loading: !!id, streaming: false, error: null, warning: null, abort: null });

// ---------------------------------------------------------------- tiny external store keyed by conversation id ("new" for drafts)
const store = new Map<string, ChatState>();
const alias = new Map<string, string>(); // "new" → id of the conversation created from the draft (so the / page keeps showing it)
const listeners = new Set<() => void>();
const key = (id: string | null) => { const k = id ?? "new"; return alias.get(k) ?? k; };
function get(id: string | null): ChatState { const k = key(id); if (!store.has(k)) store.set(k, empty(k === "new" ? null : k)); return store.get(k)!; }
function set(id: string | null, patch: Partial<ChatState> | ((s: ChatState) => Partial<ChatState>)) { const k = key(id); const cur = get(id); const p = typeof patch === "function" ? patch(cur) : patch; store.set(k, { ...cur, ...p }); listeners.forEach(l => l()); }
function subscribe(l: () => void) { listeners.add(l); return () => { listeners.delete(l); }; }

export function resetDraft() { alias.delete("new"); store.set("new", empty(null)); listeners.forEach(l => l()); }

async function load(id: string) {
  set(id, { loading: true, error: null });
  try {
    const r = await api<{ conversation: ConversationDTO; messages: MessageDTO[]; siblings: ChatState["siblings"]; attachments: AttachmentDTO[]; artifacts: ArtifactDTO[] }>(`/api/conversations/${id}?all=1`);
    set(id, { conversation: r.conversation, messages: r.messages.filter(m => m.active), allMessages: r.messages, siblings: r.siblings, attachments: r.attachments, artifacts: r.artifacts, loading: false });
  } catch (e) { set(id, { loading: false, error: (e as Error).message }); }
}

/** Streams one assistant turn; `endpoint` is one of the chat SSE routes. */
async function stream(fromId: string | null, endpoint: string, payload: Record<string, unknown>, opts: { optimisticUser?: MessageDTO; replaceAssistantId?: string } = {}) {
  const abort = new AbortController();
  let id = fromId;
  set(id, s => ({ streaming: true, error: null, warning: null, abort, messages: opts.optimisticUser ? [...s.messages, opts.optimisticUser] : s.messages }));
  let assistant: MessageDTO | null = null;
  let textBuf = ""; let reasoningBuf = ""; let raf = 0;
  const flush = () => { raf = 0; if (!assistant) return; const a = assistant; set(id, s => ({ messages: s.messages.map(m => m.id === a.id ? { ...a, parts: [...a.parts] } : m) })); };
  const schedule = () => { if (!raf) raf = requestAnimationFrame(flush); };
  try {
    const res = await fetch(endpoint, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(payload), signal: abort.signal, credentials: "same-origin" });
    for await (const { event, data } of readSSE(res)) {
      const d = data as Record<string, unknown>;
      switch (event) {
        case "conversation": {
          const newId = d.id as string;
          if (id !== newId) { // draft → real conversation: move state
            const s = get(id); store.set(newId, { ...s, conversationId: newId }); if (id === null) alias.set("new", newId);
            id = newId; listeners.forEach(l => l());
            if (d.created) { window.history.replaceState(null, "", `/c/${newId}`); void mutate("/api/conversations"); }
          }
          if (opts.optimisticUser && d.userMessageId) { const uid = d.userMessageId as string; const tmp = opts.optimisticUser.id; set(id, s => ({ messages: s.messages.map(m => m.id === tmp ? { ...m, id: uid, conversation_id: id! } : m) })); }
          if (d.regenerated || d.edited) { const old = (d.regenerated ?? d.edited) as string; set(id, s => { const idx = s.messages.findIndex(m => m.id === old); return { messages: idx >= 0 ? s.messages.slice(0, d.edited ? idx : idx) : s.messages }; }); if (d.edited && opts.optimisticUser) set(id, s => ({ messages: [...s.messages, { ...opts.optimisticUser!, id: d.userMessageId as string }] })); }
          break;
        }
        case "start": {
          assistant = { id: d.messageId as string, conversation_id: id!, parent_id: (payload.parentMessageId as string) ?? null, role: "assistant", parts: [], text: "", model_id: d.model as string, provider: d.provider as string, settings: null, routing: d.routing as MessageDTO["routing"], usage: null, cost_usd: null, latency_ms: null, ttft_ms: null, status: "streaming", error: null, active: true, created_at: new Date().toISOString() };
          textBuf = ""; reasoningBuf = "";
          set(id, s => ({ messages: [...s.messages.filter(m => m.id !== opts.replaceAssistantId), assistant!] }));
          break;
        }
        case "text": { if (!assistant) break; textBuf += d.delta as string; let tp = assistant.parts.find(p => p.type === "text") as Extract<Part, { type: "text" }> | undefined; if (!tp) { tp = { type: "text", text: "" }; assistant.parts.push(tp); } tp.text = textBuf; schedule(); break; }
        case "reasoning": { if (!assistant) break; reasoningBuf += d.delta as string; let rp = assistant.parts.find(p => p.type === "reasoning") as Extract<Part, { type: "reasoning" }> | undefined; if (!rp) { rp = { type: "reasoning", text: "", summarized: true }; assistant.parts.unshift(rp); } rp.text = reasoningBuf; schedule(); break; }
        case "tool_call": { if (!assistant) break; assistant.parts.push({ type: "tool_call", id: d.id as string, name: d.name as string, input: d.input, status: (d.status as "running") ?? "running" }); flush(); break; }
        case "tool_result": { if (!assistant) break; const tc = assistant.parts.find(p => p.type === "tool_call" && p.id === d.id) as Extract<Part, { type: "tool_call" }> | undefined; if (tc) { tc.status = d.status as "ok"; tc.output = d.output; tc.error = d.error as string | undefined; tc.durationMs = d.durationMs as number; } flush(); break; }
        case "part": { if (!assistant) break; assistant.parts.push(d as unknown as Part); flush(); break; }
        case "artifacts": { if (!assistant) break; const arts = data as { id: string; title: string; type: string; language: string }[]; assistant.parts.push(...arts.map(a => ({ type: "artifact" as const, artifactId: a.id, title: a.title, artifactType: a.type, language: a.language }))); set(id, s => ({ artifacts: [...s.artifacts, ...arts.map(a => ({ id: a.id, root_id: a.id, version: 1, type: a.type, language: a.language, title: a.title, saved: false, message_id: assistant!.id, created_at: new Date().toISOString() }))] })); flush(); break; }
        case "usage": { if (!assistant) break; assistant.usage = d.usage as MessageDTO["usage"]; assistant.cost_usd = d.costUsd as number | null; assistant.latency_ms = d.latencyMs as number; assistant.ttft_ms = d.ttftMs as number | null; assistant.status = "complete"; flush(); break; }
        case "meta": { if (assistant && d.servedBy) { assistant.routing = { ...(assistant.routing ?? {}), servedBy: d.servedBy as string }; flush(); } break; }
        case "title": { set(id, s => ({ conversation: s.conversation ? { ...s.conversation, title: d.title as string } : s.conversation })); if (!get(id).conversation && id) void load(id); void mutate("/api/conversations"); break; }
        case "warning": { set(id, { warning: d.message as string }); break; }
        case "error": { if (assistant) { assistant.status = "error"; assistant.error = d.message as string; flush(); } set(id, { error: d.message as string }); break; }
        case "done": { if (assistant && assistant.status === "streaming") { assistant.status = (d.status as string) ?? "complete"; flush(); } break; }
      }
    }
  } catch (e) {
    if ((e as Error).name !== "AbortError") { set(id, { error: (e as Error).message }); toast((e as Error).message, "error"); }
    else if (assistant) { assistant.status = "cancelled"; flush(); }
  } finally {
    if (raf) cancelAnimationFrame(raf);
    set(id, { streaming: false, abort: null });
    if (id) void load(id).then(() => mutate("/api/conversations"));
  }
}

export function useChat(conversationId: string | null) {
  const state = useSyncExternalStore(subscribe, () => get(conversationId), () => get(conversationId));
  useEffect(() => { if (conversationId && !get(conversationId).conversation && !get(conversationId).streaming) void load(conversationId); }, [conversationId]);

  const send = useCallback(async (text: string, o: ChatSendOptions = {}, atts: AttachmentDTO[] = []) => {
    const parts: Part[] = [...atts.map(a => a.mime.startsWith("image/") ? ({ type: "image" as const, attachmentId: a.id, mime: a.mime, name: a.filename }) : ({ type: "file" as const, attachmentId: a.id, mime: a.mime, name: a.filename, size: a.size ?? a.size_bytes, pages: a.pages ?? a.page_count ?? undefined })), ...(text.trim() ? [{ type: "text" as const, text }] : [])];
    const optimistic: MessageDTO = { id: `tmp_${Date.now()}`, conversation_id: conversationId ?? "", parent_id: null, role: "user", parts, text, model_id: null, provider: null, settings: null, routing: null, usage: null, cost_usd: null, latency_ms: null, ttft_ms: null, status: "complete", error: null, active: true, created_at: new Date().toISOString() };
    await stream(conversationId, "/api/chat", { conversationId: conversationId ?? undefined, text, attachmentIds: atts.map(a => a.id), ...o }, { optimisticUser: optimistic });
  }, [conversationId]);

  const regenerate = useCallback(async (messageId: string, o: Pick<ChatSendOptions, "modelId" | "routingMode" | "settings" | "toolsEnabled"> = {}) => {
    await stream(conversationId, "/api/chat/regenerate", { messageId, ...o }, { replaceAssistantId: messageId });
  }, [conversationId]);

  const edit = useCallback(async (messageId: string, text: string, o: Pick<ChatSendOptions, "modelId" | "routingMode" | "toolsEnabled"> = {}) => {
    const old = get(conversationId).messages.find(m => m.id === messageId);
    const optimistic: MessageDTO = { ...(old as MessageDTO), id: `tmp_${Date.now()}`, parts: [...(old?.parts.filter(p => p.type !== "text") ?? []), { type: "text", text }], text };
    await stream(conversationId, "/api/chat/edit", { messageId, text, ...o }, { optimisticUser: optimistic });
  }, [conversationId]);

  const approve = useCallback(async (messageId: string, approveIds: string[], denyIds: string[] = []) => {
    await stream(conversationId, "/api/chat/approve", { messageId, approve: approveIds, deny: denyIds });
  }, [conversationId]);

  const stop = useCallback(() => { get(conversationId).abort?.abort(); }, [conversationId]);
  const reload = useCallback(() => { if (conversationId) void load(conversationId); }, [conversationId]);
  const switchBranch = useCallback(async (targetId: string) => { if (!conversationId) return; await api(`/api/conversations/${conversationId}/branch`, { method: "POST", json: { activate: targetId } }); await load(conversationId); }, [conversationId]);

  return { ...state, send, regenerate, edit, approve, stop, reload, switchBranch };
}