freezmo/components/chat/ChatScreen.tsx

"use client";

import { useCallback, useEffect, useMemo, useRef, useState, useSyncExternalStore } from "react";
import { useRouter, useSearchParams } from "next/navigation";
import { useReducedMotion } from "framer-motion";
import Link from "next/link";

import { Composer, type PendingAttachment } from "./Composer";
import { ContextPicker } from "./ContextPicker";
import { MessageRow } from "./MessageRow";
import { ModelPicker } from "./ModelPicker";
import { ChatIcon, ChevronDownIcon, ChevronLeftIcon, PlusIcon, RefreshIcon } from "@/components/icons";
import { ModelErrorActions } from "@/components/ModelErrorActions";
import { ThemeToggle } from "@/components/ThemeToggle";
import { Button } from "@/components/ui";
import { ApiError } from "@/lib/client-api";
import { extractChatAttachment, streamChat } from "@/lib/chat-client";
import { buildChatContext, deckHasNotes } from "@/lib/chat-context";
import { takeChatSeed, type ChatSeed } from "@/lib/chat-seed";
import {
  chatMessageId,
  readStoredChatModel,
  storeChatModel,
  type ChatAttachment,
  type ChatMessage,
  type ChatModelChoice,
  type ChatScope,
} from "@/lib/chat";
import { extractDocument } from "@/lib/documents";
import { folderStore } from "@/lib/folders";
import { isImageFile, prepareImage } from "@/lib/image";
import { useI18n } from "@/lib/i18n";
import { showToast } from "@/lib/toast";
import { deckStore } from "@/lib/storage";
import { settingsStore } from "@/lib/settings";
import { useHealth } from "@/lib/use-health";
import { useOnline } from "@/lib/use-online";

type ExtraChunk = { deckId: string; deckTitle: string; section: string; text: string };

/**
 * Chat mode: an LLM-style conversation that knows the user's notes. The
 * retrieval runs here in the browser (BM25 over the scoped decks), and only
 * the top chunks plus the recent turns travel with each request.
 */
export function ChatScreen() {
  const { t, plural } = useI18n();
  const online = useOnline();
  const reduceMotion = useReducedMotion();
  const router = useRouter();
  const search = useSearchParams();

  const decks = useSyncExternalStore(deckStore.subscribe, deckStore.getSnapshot, deckStore.getServerSnapshot);
  const folders = useSyncExternalStore(folderStore.subscribe, folderStore.getSnapshot, folderStore.getServerSnapshot);
  const settings = useSyncExternalStore(settingsStore.subscribe, settingsStore.getSnapshot, settingsStore.getServerSnapshot);
  const { health } = useHealth();

  const deckParam = search.get("deck");

  const [scope, setScope] = useState<ChatScope>(() =>
    deckParam ? { kind: "decks", deckIds: [deckParam] } : { kind: "all" },
  );
  const [seed, setSeed] = useState<ChatSeed | null>(null);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [draft, setDraft] = useState("");
  const [attachments, setAttachments] = useState<PendingAttachment[]>([]);
  const [attachBusy, setAttachBusy] = useState(false);
  const [model, setModel] = useState("");
  const [busy, setBusy] = useState(false);
  const [pinned, setPinned] = useState(true);
  const [scopeOpen, setScopeOpen] = useState(false);
  const [modelOpen, setModelOpen] = useState(false);
  const [error, setError] = useState<{ message: string; kind?: string } | null>(null);

  const abortRef = useRef<AbortController | null>(null);
  const scrollRef = useRef<HTMLDivElement>(null);
  const pinnedRef = useRef(true);
  const extrasRef = useRef<Record<string, ExtraChunk[]>>({});

  /* Hand-offs from a deck menu / results screen. */
  useEffect(() => {
    const taken = takeChatSeed();
    if (!taken) return;
    setSeed(taken);
    if (taken.deckId) setScope({ kind: "decks", deckIds: [taken.deckId] });
  }, []);

  useEffect(() => {
    setModel(readStoredChatModel() ?? "");
  }, []);

  /* Keyboard: keep the composer above the on-screen keyboard. */
  useEffect(() => {
    const viewport = window.visualViewport;
    if (!viewport) return;
    const apply = () => {
      const inset = Math.max(0, window.innerHeight - viewport.height - viewport.offsetTop);
      document.documentElement.style.setProperty("--chat-keyboard", `${Math.round(inset)}px`);
    };
    apply();
    viewport.addEventListener("resize", apply);
    viewport.addEventListener("scroll", apply);
    return () => {
      viewport.removeEventListener("resize", apply);
      viewport.removeEventListener("scroll", apply);
      document.documentElement.style.removeProperty("--chat-keyboard");
    };
  }, []);

  const chatModels = useMemo<ChatModelChoice[]>(() => {
    const fromHealth = health?.chatModels ?? [];
    const list =
      fromHealth.length > 0
        ? fromHealth
        : [
            { name: health?.text?.configured ?? "", badge: "text" as const, available: health?.text?.available ?? null },
            { name: health?.vision?.configured ?? "", badge: "vision" as const, available: health?.vision?.available ?? null },
          ];
    return list.filter((entry) => entry.name);
  }, [health]);

  const activeModel = useMemo(() => {
    if (model && chatModels.some((entry) => entry.name === model)) return model;
    return chatModels[0]?.name ?? model;
  }, [model, chatModels]);

  const scopeText = useMemo(() => {
    if (scope.kind === "all") return t("chat.scope.all");
    if (scope.kind === "none") return t("chat.scope.none");
    return plural(scope.deckIds.length, "chat.scope.decks");
  }, [scope, t, plural]);

  const scopedDecks = useMemo(() => {
    if (scope.kind === "none") return [];
    if (scope.kind === "all") return decks;
    return decks.filter((deck) => scope.deckIds.includes(deck.id));
  }, [decks, scope]);

  const missingNotes = useMemo(
    () => (settings.chatNotes ? scopedDecks.some((deck) => !deckHasNotes(deck)) : false),
    [scopedDecks, settings.chatNotes],
  );

  const suggestions = useMemo(() => {
    const chosen = scope.kind === "decks" ? scopedDecks : [];
    if (chosen.length === 1) {
      return [
        t("chat.suggest.explain", { deck: chosen[0].title }),
        t("chat.suggest.quiz", { deck: chosen[0].title }),
        t("chat.suggest.summarize"),
      ];
    }
    if (scope.kind === "none") return [t("chat.suggest.general"), t("chat.suggest.explainTopic")];
    return [t("chat.suggest.summarize"), t("chat.suggest.quizAll"), t("chat.suggest.studyFirst")];
  }, [scope, scopedDecks, t]);

  const scrollToBottom = useCallback(
    (smooth = false) => {
      const element = scrollRef.current;
      if (!element) return;
      element.scrollTo({ top: element.scrollHeight, behavior: smooth && !reduceMotion ? "smooth" : "auto" });
    },
    [reduceMotion],
  );

  useEffect(() => {
    if (pinnedRef.current) scrollToBottom(false);
  }, [messages, scrollToBottom]);

  const onScroll = () => {
    const element = scrollRef.current;
    if (!element) return;
    const distance = element.scrollHeight - element.scrollTop - element.clientHeight;
    const next = distance < 80;
    pinnedRef.current = next;
    setPinned(next);
  };

  const extraFor = (text: string, pending: PendingAttachment[]): ExtraChunk[] => {
    const chunks: ExtraChunk[] = [];
    if (seed) chunks.push({ deckId: seed.deckId ?? "seed", deckTitle: seed.label, section: seed.label, text: seed.text });
    if (pending.length > 0) {
      chunks.push({
        deckId: "attachment",
        deckTitle: t("chat.attachmentLabel"),
        section: t("chat.attachmentLabel"),
        text: pending.map((entry) => `${entry.attachment.name}\n${entry.text}`).join("\n\n").slice(0, 24_000),
      });
    }
    void text;
    return chunks;
  };

  const run = async (history: ChatMessage[], extraChunks: ExtraChunk[]) => {
    const lastUser = history[history.length - 1];
    const context = buildChatContext({
      decks,
      scope,
      question: lastUser.content,
      history: history.slice(0, -1),
      useNotes: settings.chatNotes,
      extraChunks,
    });
    const sources = context.chunks.map(({ deckId, deckTitle, section }) => ({ deckId, deckTitle, section }));
    const assistantId = chatMessageId("a");

    setMessages([...history, { id: assistantId, role: "assistant", content: "", sources, at: Date.now() }]);
    setBusy(true);
    pinnedRef.current = true;
    setPinned(true);
    const controller = new AbortController();
    abortRef.current = controller;
    let produced = false;

    try {
      await streamChat({
        model: activeModel || undefined,
        messages: history.map((message) => ({ role: message.role, content: message.content })),
        context: context.chunks.map(({ deckId, deckTitle, section, text }) => ({ deckId, deckTitle, section, text })),
        signal: controller.signal,
        onToken: (text) => {
          produced = true;
          setMessages((previous) =>
            previous.map((message) =>
              message.id === assistantId ? { ...message, content: message.content + text } : message,
            ),
          );
        },
        onThinking: (text) => {
          setMessages((previous) =>
            previous.map((message) =>
              message.id === assistantId ? { ...message, thinking: (message.thinking ?? "") + text } : message,
            ),
          );
        },
        onDone: (modelName) => {
          setMessages((previous) =>
            previous.map((message) => (message.id === assistantId ? { ...message, model: modelName } : message)),
          );
        },
      });
    } catch (error) {
      const apiError = error instanceof ApiError ? error : new ApiError(String(error), "unknown");
      if (apiError.kind === "aborted") {
        // Stopped by the user: keep whatever streamed in.
      } else {
        setError({ message: apiError.message, kind: apiError.kind });
        if (!produced) {
          // Nothing arrived: take the placeholder away and give the text back.
          setMessages(history);
          setDraft(lastUser.content);
        } else {
          setMessages((previous) =>
            previous.map((message) =>
              message.id === assistantId
                ? { ...message, error: { message: apiError.message, kind: apiError.kind } }
                : message,
            ),
          );
        }
      }
    } finally {
      abortRef.current = null;
      setBusy(false);
    }
  };

  const send = (text?: string) => {
    const content = (text ?? draft).trim();
    if (!content || busy || !online) return;
    const pending = attachments;
    const userMessage: ChatMessage = {
      id: chatMessageId("u"),
      role: "user",
      content,
      attachments: pending.length > 0 ? pending.map((entry) => entry.attachment) : undefined,
      at: Date.now(),
    };
    const extras = extraFor(content, pending);
    extrasRef.current[userMessage.id] = extras;
    const history = [...messages, userMessage];
    setError(null);
    setMessages(history);
    setDraft("");
    setAttachments([]);
    void run(history, extras);
  };

  const retryLast = () => {
    const lastUserIndex = messages.map((message) => message.role).lastIndexOf("user");
    if (lastUserIndex < 0) return;
    const history = messages.slice(0, lastUserIndex + 1);
    setMessages(history);
    void run(history, extrasRef.current[history[lastUserIndex].id] ?? []);
  };

  const regenerate = () => {
    const lastUserIndex = messages.map((message) => message.role).lastIndexOf("user");
    if (lastUserIndex < 0) return;
    const history = messages.slice(0, lastUserIndex + 1);
    setMessages(history);
    void run(history, extrasRef.current[history[lastUserIndex].id] ?? []);
  };

  const editLast = () => {
    const lastUserIndex = messages.map((message) => message.role).lastIndexOf("user");
    if (lastUserIndex < 0) return;
    setDraft(messages[lastUserIndex].content);
    setMessages(messages.slice(0, lastUserIndex));
  };

  const newChat = () => {
    abortRef.current?.abort();
    setMessages([]);
    setSeed(null);
    setAttachments([]);
    setDraft("");
  };

  const handleAttach = async (file: File) => {
    setAttachBusy(true);
    try {
      if (isImageFile(file)) {
        const prepared = await prepareImage(file);
        const text = await extractChatAttachment({ images: [prepared.dataUrl] });
        const attachment: ChatAttachment = { name: file.name || t("chat.attachPhoto"), kind: "photo", chars: text.length };
        setAttachments((previous) => [...previous, { attachment, text }]);
      } else {
        const document = await extractDocument(file);
        if (!document.text.trim()) {
          showToast({ message: t("chat.attachNoText") });
        } else {
          const attachment: ChatAttachment = { name: document.name, kind: "document", chars: document.text.length };
          setAttachments((previous) => [...previous, { attachment, text: document.text }]);
        }
      }
    } catch (error) {
      showToast({ message: error instanceof ApiError ? error.message : t("chat.attachFailed") });
    } finally {
      setAttachBusy(false);
    }
  };

  const lastAssistantId = [...messages].reverse().find((message) => message.role === "assistant")?.id;
  const lastUserId = [...messages].reverse().find((message) => message.role === "user")?.id;
  const lastFailedAssistant = [...messages].reverse().find((message) => message.role === "assistant" && message.error)?.id;

  return (
    <div className="relative flex h-full min-h-0 flex-col" style={{ paddingBottom: "var(--chat-keyboard, 0px)" }}>
      <header className="glass sticky top-0 z-30 flex h-14 shrink-0 items-center gap-1.5 px-2 sm:px-4">
        <button
          type="button"
          onClick={() => (window.history.length > 1 ? router.back() : router.push("/"))}
          aria-label={t("chat.back")}
          title={t("chat.back")}
          className="grid size-10 shrink-0 place-items-center rounded-2xl text-[var(--text-muted)] transition hover:bg-[var(--hover)] hover:text-[var(--text)] lg:hidden"
        >
          <ChevronLeftIcon className="size-5" />
        </button>
        <div className="min-w-0 flex-1">
          <h1 className="truncate text-[15px] font-semibold text-[var(--text)]">{t("chat.title")}</h1>
          <p className="truncate text-[11.5px] text-[var(--text-subtle)]">
            {scopeText}
            {seed ? ` · ${seed.label}` : ""}
          </p>
        </div>
        <ThemeToggle className="hidden lg:grid" />
        <button
          type="button"
          onClick={newChat}
          aria-label={t("chat.new")}
          title={t("chat.new")}
          className="flex h-10 shrink-0 items-center gap-1.5 rounded-2xl border border-[var(--glass-border)] bg-[var(--surface-strong)] px-3 text-[12.5px] font-medium text-[var(--text-muted)] transition hover:bg-[var(--hover)] hover:text-[var(--text)]"
        >
          <PlusIcon className="size-4" />
          <span className="hidden sm:inline">{t("chat.new")}</span>
        </button>
      </header>

      <div className="relative min-h-0 flex-1">
        <div ref={scrollRef} onScroll={onScroll} className="h-full overflow-y-auto overscroll-contain">
          <div className="mx-auto w-full max-w-3xl space-y-4 px-4 py-4">
            {messages.length === 0 ? (
              <div className="space-y-5 py-6">
                <div className="text-center">
                  <ChatIcon className="mx-auto size-10 text-accent-400" />
                  <h2 className="mt-3 text-[20px] font-semibold tracking-[-0.01em] text-[var(--text)]">{t("chat.emptyTitle")}</h2>
                  <p className="mx-auto mt-1.5 max-w-sm text-[13.5px] leading-relaxed text-[var(--text-subtle)]">
                    {t("chat.emptyHint")}
                  </p>
                </div>
                {missingNotes ? (
                  <p className="rounded-2xl border border-amber-300/50 bg-amber-50/70 px-3.5 py-3 text-[12.5px] leading-relaxed text-amber-800 dark:border-amber-400/25 dark:bg-amber-400/10 dark:text-amber-200">
                    {t("chat.oldDecksNote")}
                  </p>
                ) : null}
                <div className="flex flex-wrap justify-center gap-2">
                  {suggestions.map((suggestion) => (
                    <button
                      key={suggestion}
                      type="button"
                      disabled={!online || busy}
                      onClick={() => send(suggestion)}
                      className="max-w-full rounded-2xl border border-[var(--glass-border)] bg-[var(--surface-strong)] px-3.5 py-2.5 text-left text-[13px] text-[var(--text-muted)] transition hover:bg-[var(--hover)] hover:text-[var(--text)] disabled:opacity-55"
                    >
                      {suggestion}
                    </button>
                  ))}
                </div>
              </div>
            ) : (
              <>
                {seed ? (
                  <div className="flex items-center gap-2 rounded-2xl border border-[var(--hairline-soft)] bg-[var(--surface-sunken)] px-3 py-2 text-[12px] text-[var(--text-subtle)]">
                    <span className="truncate">
                      {t("chat.fromLabel")}: {seed.label}
                    </span>
                  </div>
                ) : null}
                {messages.map((message) => (
                  <MessageRow
                    key={message.id}
                    message={message}
                    streaming={busy && message.id === lastAssistantId}
                    isLastAssistant={message.id === lastAssistantId}
                    isLastUser={message.id === lastUserId}
                    busy={busy}
                    onRegenerate={regenerate}
                    onEdit={editLast}
                    onRetry={message.id === lastFailedAssistant ? retryLast : undefined}
                  />
                ))}
              </>
            )}
            <div className="h-2" aria-hidden />
          </div>
        </div>

        {!pinned && messages.length > 0 ? (
          <button
            type="button"
            onClick={() => {
              pinnedRef.current = true;
              setPinned(true);
              scrollToBottom(true);
            }}
            aria-label={t("chat.jumpToBottom")}
            title={t("chat.jumpToBottom")}
            className="glass-panel absolute bottom-4 left-1/2 grid size-10 -translate-x-1/2 place-items-center rounded-full text-[var(--text-muted)] shadow-lg"
          >
            <ChevronDownIcon className="size-5" />
          </button>
        ) : null}
      </div>

      {error ? (
        <div
          role="alert"
          className="mx-3 mb-2 shrink-0 rounded-2xl border border-rose-300/60 bg-rose-50/80 p-3 dark:border-rose-400/30 dark:bg-rose-500/10"
        >
          <p className="text-[13px] leading-relaxed text-rose-800 dark:text-rose-200">{error.message}</p>
          <div className="mt-2 flex flex-wrap items-center gap-2">
            {messages.some((message) => message.role === "user") ? (
              <Button size="sm" variant="secondary" onClick={retryLast}>
                <RefreshIcon className="size-4" /> {t("chat.retry")}
              </Button>
            ) : null}
            <ModelErrorActions kind={error.kind} />
          </div>
        </div>
      ) : null}

      <div className="relative shrink-0">
        <Composer
          value={draft}
          onChange={setDraft}
          onSend={() => send()}
          onStop={() => abortRef.current?.abort()}
          busy={busy}
          online={online}
          scopeLabel={scopeText}
          modelLabel={activeModel || t("chat.modelTitle")}
          onOpenScope={() => setScopeOpen(true)}
          onOpenModel={() => setModelOpen(true)}
          onAttachFile={(file) => void handleAttach(file)}
          attachBusy={attachBusy}
          attachments={attachments}
          onRemoveAttachment={(index) => setAttachments((previous) => previous.filter((_, i) => i !== index))}
        />
        <ModelPicker
          open={modelOpen}
          onClose={() => setModelOpen(false)}
          models={chatModels}
          value={activeModel}
          onChange={(next) => {
            setModel(next);
            storeChatModel(next);
          }}
        />
        <ContextPicker
          open={scopeOpen}
          onClose={() => setScopeOpen(false)}
          scope={scope}
          onChange={setScope}
          decks={decks}
          folders={folders}
        />
      </div>

      {settings.chatNotes === false && scope.kind !== "none" ? (
        <p className="border-t border-[var(--hairline-soft)] bg-[var(--surface-sunken)] px-4 py-2 text-[11.5px] text-[var(--text-subtle)]">
          {t("chat.notesOffNote")}{" "}
          <Link href="/" className="underline decoration-dotted">
            {t("chat.notesOffLink")}
          </Link>
        </p>
      ) : null}
    </div>
  );
}