// Author: Simon-Pierre Boucher
// Contact: contact@spboucher.ai
// Project: chat.spboucher.ai

"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  type ApiConversation,
  type ApiMessage,
  type ApiModel,
  computeThread,
  choicesForLeaf,
  estimateTokensClient,
} from "./types";
import { Sidebar } from "./Sidebar";
import { MessageList } from "./MessageList";
import { Composer } from "./Composer";
import { ModelSheet } from "./ModelSheet";

interface StreamingState {
  assistantMessageId: string;
  generationId: string;
  content: string;
  reasoning: string;
  error: string | null;
}

export function ChatApp() {
  const [conversations, setConversations] = useState<ApiConversation[]>([]);
  const [activeId, setActiveId] = useState<string | null>(null);
  const [messages, setMessages] = useState<ApiMessage[]>([]);
  const [branchChoice, setBranchChoice] = useState<Map<string, string>>(new Map());
  const [models, setModels] = useState<ApiModel[]>([]);
  const [modelId, setModelId] = useState<string | null>(null);
  const [sheetOpen, setSheetOpen] = useState(false);
  const [drawerOpen, setDrawerOpen] = useState(false);
  const [streaming, setStreaming] = useState<StreamingState | null>(null);
  const [pendingUser, setPendingUser] = useState<ApiMessage | null>(null);
  const abortRef = useRef<AbortController | null>(null);
  const activeIdRef = useRef<string | null>(null);
  activeIdRef.current = activeId;

  const model = useMemo(() => models.find((m) => m.id === modelId) ?? null, [models, modelId]);

  const refreshConversations = useCallback(async () => {
    const res = await fetch("/api/conversations");
    if (!res.ok) return;
    const json = await res.json();
    setConversations(json.conversations);
  }, []);

  const loadConversation = useCallback(async (id: string) => {
    const res = await fetch(`/api/conversations/${id}`);
    if (!res.ok) return;
    const json = await res.json();
    if (activeIdRef.current !== id) return; // user already moved on
    setMessages(json.messages);
    const leaf = json.conversation.current_leaf_id;
    setBranchChoice(leaf ? choicesForLeaf(json.messages, leaf) : new Map());
  }, []);

  // Initial load: conversations + model catalog.
  useEffect(() => {
    refreshConversations();
    (async () => {
      const res = await fetch("/api/models");
      if (!res.ok) return;
      const json = await res.json();
      const all: ApiModel[] = json.models;
      setModels(all);
      const saved = localStorage.getItem("spb-model");
      const avail = all.filter((m) => m.available);
      const pick =
        (saved && avail.find((m) => m.id === saved)) ||
        avail.find((m) => m.favorite) ||
        avail.filter((m) => m.lastUsedAt).sort((a, b) => (b.lastUsedAt ?? 0) - (a.lastUsedAt ?? 0))[0] ||
        avail.find((m) => m.id === "anthropic/claude-sonnet-4.5") ||
        avail.find((m) => m.id === "openai/gpt-4o-mini") ||
        avail[0];
      if (pick) setModelId(pick.id);
    })();
  }, [refreshConversations]);

  useEffect(() => {
    if (modelId) localStorage.setItem("spb-model", modelId);
  }, [modelId]);

  useEffect(() => {
    if (activeId) loadConversation(activeId);
    else {
      setMessages([]);
      setBranchChoice(new Map());
    }
  }, [activeId, loadConversation]);

  const thread = useMemo(() => computeThread(messages, branchChoice), [messages, branchChoice]);

  const contextTokens = useMemo(
    () => thread.reduce((acc, m) => acc + estimateTokensClient(m.content), 0),
    [thread]
  );

  /** Core streaming loop shared by send + regenerate. */
  const runStream = useCallback(
    async (body: Record<string, unknown>) => {
      const ctrl = new AbortController();
      abortRef.current = ctrl;
      let convId = (body.conversationId as string) ?? null;
      try {
        const res = await fetch("/api/chat", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
          signal: ctrl.signal,
        });
        if (!res.ok || !res.body) {
          const json = await res.json().catch(() => ({}));
          setStreaming((s) =>
            s ? { ...s, error: json.error ?? "The request failed." } : {
              assistantMessageId: "",
              generationId: "",
              content: "",
              reasoning: "",
              error: json.error ?? "The request failed.",
            }
          );
          return;
        }

        convId = res.headers.get("X-Conversation-Id") ?? convId;
        if (convId && activeIdRef.current !== convId) {
          setActiveId(convId);
          refreshConversations();
        }

        const reader = res.body.getReader();
        const decoder = new TextDecoder();
        let buffer = "";
        for (;;) {
          const { done, value } = await reader.read();
          if (done) break;
          buffer += decoder.decode(value, { stream: true });
          let idx: number;
          while ((idx = buffer.indexOf("\n")) !== -1) {
            const line = buffer.slice(0, idx).trim();
            buffer = buffer.slice(idx + 1);
            if (!line.startsWith("data:")) continue;
            const payload = line.slice(5).trim();
            if (!payload) continue;
            let event: Record<string, unknown>;
            try {
              event = JSON.parse(payload);
            } catch {
              continue;
            }
            switch (event.type) {
              case "meta":
                setStreaming({
                  assistantMessageId: event.assistantMessageId as string,
                  generationId: event.generationId as string,
                  content: "",
                  reasoning: "",
                  error: null,
                });
                break;
              case "content.delta":
                setStreaming((s) => (s ? { ...s, content: s.content + (event.text as string) } : s));
                break;
              case "reasoning.delta":
                setStreaming((s) => (s ? { ...s, reasoning: s.reasoning + (event.text as string) } : s));
                break;
              case "generation.error":
                setStreaming((s) => (s ? { ...s, error: event.message as string } : s));
                break;
            }
          }
        }
      } catch {
        // network drop or user abort — server state is canonical; resync below
      } finally {
        abortRef.current = null;
        setPendingUser(null);
        setStreaming(null);
        const target = convId ?? activeIdRef.current;
        if (target) {
          if (activeIdRef.current === null) setActiveId(target);
          await loadConversation(target);
        }
        refreshConversations();
      }
    },
    [loadConversation, refreshConversations]
  );

  const sendMessage = useCallback(
    async (content: string) => {
      if (!modelId || streaming) return;
      const parentId = thread.length ? thread[thread.length - 1].id : null;
      // Optimistic: show the user message instantly, marked pending.
      setPendingUser({
        id: `pending-${Date.now()}`,
        conversation_id: activeId ?? "",
        parent_id: parentId,
        role: "user",
        content,
        reasoning: null,
        model_id: null,
        model_name: null,
        provider: null,
        generation_id: null,
        status: "pending",
        error_message: null,
        created_at: Date.now(),
      });
      await runStream({
        conversationId: activeId ?? undefined,
        parentId,
        content,
        modelId,
      });
    },
    [modelId, streaming, thread, activeId, runStream]
  );

  const regenerate = useCallback(
    async (assistantMessageId: string, withModelId?: string) => {
      if (streaming) return;
      await runStream({
        regenerateOf: assistantMessageId,
        modelId: withModelId ?? modelId,
      });
    },
    [streaming, modelId, runStream]
  );

  const stopGeneration = useCallback(async () => {
    const genId = streaming?.generationId;
    if (genId) {
      await fetch(`/api/generations/${genId}/cancel`, { method: "POST" }).catch(() => {});
    }
    abortRef.current?.abort();
  }, [streaming]);

  const selectBranch = useCallback(
    (parentKey: string, childId: string) => {
      const next = new Map(branchChoice);
      next.set(parentKey, childId);
      setBranchChoice(next);
      // Persist the new active leaf so other devices resume the same branch.
      const newThread = computeThread(messages, next);
      const leaf = newThread[newThread.length - 1];
      if (leaf && activeId) {
        fetch(`/api/conversations/${activeId}`, {
          method: "PATCH",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ currentLeafId: leaf.id }),
        }).catch(() => {});
      }
    },
    [branchChoice, messages, activeId]
  );

  const newConversation = useCallback(() => {
    setActiveId(null);
    setDrawerOpen(false);
  }, []);

  const deleteConversation = useCallback(
    async (id: string) => {
      await fetch(`/api/conversations/${id}`, { method: "DELETE" });
      if (activeIdRef.current === id) setActiveId(null);
      refreshConversations();
    },
    [refreshConversations]
  );

  const toggleFavorite = useCallback(async (id: string, favorite: boolean) => {
    setModels((ms) => ms.map((m) => (m.id === id ? { ...m, favorite } : m)));
    await fetch("/api/models/prefs", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ modelId: id, favorite }),
    }).catch(() => {});
  }, []);

  const activeConv = conversations.find((c) => c.id === activeId) ?? null;

  return (
    <div className="app">
      <Sidebar
        conversations={conversations}
        activeId={activeId}
        open={drawerOpen}
        onSelect={(id) => {
          setActiveId(id);
          setDrawerOpen(false);
        }}
        onNew={newConversation}
        onDelete={deleteConversation}
        onClose={() => setDrawerOpen(false)}
      />
      <div className="main">
        <header className="topbar">
          <button
            className="icon-btn menu-btn"
            aria-label="Open conversations"
            onClick={() => setDrawerOpen(true)}
          >
            <svg width="20" height="20" viewBox="0 0 20 20" fill="none" stroke="currentColor" strokeWidth="1.6">
              <path d="M3 5h14M3 10h14M3 15h14" strokeLinecap="round" />
            </svg>
          </button>
          <span className="topbar-title">{activeConv?.title ?? "New conversation"}</span>
          <button className="icon-btn" aria-label="New conversation" onClick={newConversation}>
            <svg width="20" height="20" viewBox="0 0 20 20" fill="none" stroke="currentColor" strokeWidth="1.6">
              <path d="M10 4v12M4 10h12" strokeLinecap="round" />
            </svg>
          </button>
        </header>

        <MessageList
          thread={thread}
          allMessages={messages}
          pendingUser={pendingUser}
          streaming={streaming}
          onSelectBranch={selectBranch}
          onRegenerate={regenerate}
          models={models}
          currentModelName={model?.name ?? null}
        />

        <Composer
          model={model}
          contextTokens={contextTokens}
          streaming={Boolean(streaming)}
          onSend={sendMessage}
          onStop={stopGeneration}
          onOpenModelSheet={() => setSheetOpen(true)}
        />
      </div>

      {sheetOpen && (
        <ModelSheet
          models={models}
          selectedId={modelId}
          onSelect={(id) => {
            setModelId(id);
            setSheetOpen(false);
          }}
          onToggleFavorite={toggleFavorite}
          onClose={() => setSheetOpen(false)}
        />
      )}
    </div>
  );
}
