Document Q&A

"use client";

import * as React from "react";
import { ArrowUpIcon, DocumentTextIcon, SparklesIcon, StopIcon } from "@heroicons-animated/react";

import { cn } from "@/lib/utils";
import { Button } from "@/components/ui/button";
import { CitedText, type Source } from "@/components/pandacoderz-ui/citation";
import { FileDrop, FileItem, FileList, type DroppedFile } from "@/components/pandacoderz-ui/file-drop";
import { Message, MessageAvatar, MessageContent, MessageStack } from "@/components/pandacoderz-ui/message";
import { PromptInput, PromptInputAction, PromptInputActionGroup, PromptInputActions, PromptInputTextarea } from "@/components/pandacoderz-ui/prompt-input";
import { Suggestion, SuggestionList, Suggestions } from "@/components/pandacoderz-ui/suggestions";
import { TextShimmer } from "@/components/pandacoderz-ui/text-shimmer";
import { Thread, ThreadContent, ThreadScrollToBottom } from "@/components/pandacoderz-ui/thread";
import { sleep, uid } from "@/lib/mock-stream";
import { useStreamText } from "@/lib/use-stream-text";

type Doc = DroppedFile & { pages: string[] };

const docs: Doc[] = [
  {
    id: "d1", name: "Q3-board-update.pdf", size: 1_240_000, type: "application/pdf",
    pages: [
      "Q3 Board Update. Revenue grew 18% quarter over quarter to $4.2M ARR, driven by expansion in mid-market accounts.",
      "Product. We shipped the agent runtime and streaming chat. Weekly active workspaces rose from 1,900 to 2,600.",
      "Hiring. Headcount is 42, up from 36. Two open roles remain in platform engineering and one in design.",
      "Risks. Gross margin dipped to 71% on inference costs. We expect a recovery to 76% in Q4 as caching lands.",
      "Ask. We are requesting approval for a $1.5M extension of the infrastructure budget to cover the Q4 caching rollout.",
    ],
  },
  { id: "d2", name: "vendor-contract.pdf", size: 382_000, type: "application/pdf", pages: ["Master services agreement between Northwind and Acme Cloud.", "Term: 24 months with automatic renewal unless cancelled 60 days prior.", "Fees: $12,000 per month, invoiced quarterly in advance."] },
];

type Turn = { id: string; role: "user" | "assistant"; text: string; pages?: number[] };

const canned: { match: RegExp; text: string; pages: number[] }[] = [
  { match: /revenue|arr|grow/i, text: "Revenue grew 18% quarter over quarter to $4.2M ARR, with expansion in mid-market accounts as the main driver [1].", pages: [1] },
  { match: /margin|risk|cost/i, text: "Gross margin dipped to 71% because of inference costs [1]. The update expects it to recover to 76% in Q4 once caching lands [1], which is also what the $1.5M infrastructure ask is meant to fund [2].", pages: [4, 5] },
  { match: /hir|headcount|team/i, text: "Headcount is 42, up from 36 last quarter [1]. Three roles remain open: two in platform engineering and one in design [1].", pages: [3] },
  { match: /ask|approv|budget/i, text: "The board is being asked to approve a $1.5M extension of the infrastructure budget for the Q4 caching rollout [1]. The rationale is the margin dip described earlier [2].", pages: [5, 4] },
  { match: /.*/, text: "This is the Q3 board update. Revenue is up 18% to $4.2M ARR [1], the team shipped the agent runtime and streaming chat [2], and the main risk is a gross margin dip to 71% on inference costs [3].", pages: [1, 2, 4] },
];

const suggestions = ["What happened to revenue?", "Why did margin dip?", "What is the ask?"];

export type DocumentQAProps = { className?: string };

export default function DocumentQA({ className }: DocumentQAProps) {
  const [files, setFiles] = React.useState<Doc[]>(docs);
  const [selectedId, setSelectedId] = React.useState(docs[0].id);
  const [activePage, setActivePage] = React.useState<number | null>(null);
  const [turns, setTurns] = React.useState<Turn[]>([]);
  const [draft, setDraft] = React.useState("");
  const [pending, setPending] = React.useState<number[] | null>(null);
  const { text, isStreaming, start, stop, reset } = useStreamText(12);
  const pageRefs = React.useRef<Record<number, HTMLElement | null>>({});

  const doc = files.find((f) => f.id === selectedId) ?? files[0];

  const sourcesFor = (pages: number[]): Source[] =>
    pages.map((p) => ({ id: `${doc.id}-p${p}`, title: `Page ${p}`, url: `#page-${p}`, site: doc.name, snippet: doc.pages[p - 1] }));

  const focusPage = (p: number) => {
    setActivePage(p);
    pageRefs.current[p]?.scrollIntoView({ block: "nearest", behavior: "smooth" });
  };

  const ask = async (q: string) => {
    const v = q.trim();
    if (!v || isStreaming) return;
    setDraft("");
    const found = canned.find((c) => c.match.test(v)) ?? canned.at(-1)!;
    setTurns((prev) => [...prev, { id: uid("u"), role: "user", text: v }]);
    setPending(found.pages);
    reset();
    await sleep(600);
    focusPage(found.pages[0]);
    await start(found.text, () => {
      setTurns((prev) => [...prev, { id: uid("a"), role: "assistant", text: found.text, pages: found.pages }]);
      setPending(null);
    });
  };

  const addFiles = (incoming: File[]) => {
    const next: Doc[] = incoming.map((f) => ({ id: uid("f"), name: f.name, size: f.size, type: f.type, progress: 0, pages: ["(Indexed content would appear here.)"] }));
    setFiles((prev) => [...prev, ...next]);
    next.forEach((f) => {
      let p = 0;
      const id = setInterval(() => {
        p += 25;
        setFiles((prev) => prev.map((x) => (x.id === f.id ? { ...x, progress: Math.min(100, p) } : x)));
        if (p >= 100) clearInterval(id);
      }, 200);
    });
  };

  const onCitationClick = (e: React.MouseEvent) => {
    const a = (e.target as HTMLElement).closest<HTMLAnchorElement>("[data-slot=citation]");
    if (!a) return;
    e.preventDefault();
    const m = /#page-(\d+)/.exec(a.getAttribute("href") ?? "");
    if (m) focusPage(Number(m[1]));
  };

  const activeSourceId = activePage ? `${doc.id}-p${activePage}` : null;

  return (
    <div data-slot="document-qa" className={cn("grid h-full min-h-0 w-full overflow-hidden rounded-3xl border bg-background shadow-xs lg:grid-cols-[minmax(0,15rem)_minmax(0,17rem)_minmax(0,1fr)]", className)}>
      <aside className="flex min-h-0 flex-col gap-3 border-b bg-surface p-3 lg:border-r lg:border-b-0">
        <FileDrop onFiles={addFiles} accept=".pdf,.md,.docx" compact />
        <FileList className="min-h-0 flex-1 overflow-y-auto">
          {files.map((f) => (
            <FileItem key={f.id} file={f} selected={f.id === selectedId} onClick={() => { setSelectedId(f.id); setActivePage(null); }} className="cursor-pointer" onRemove={files.length > 1 ? (id) => setFiles((prev) => { const next = prev.filter((x) => x.id !== id); if (id === selectedId) setSelectedId(next[0].id); return next; }) : undefined} />
          ))}
        </FileList>
      </aside>

      <section className="hidden min-h-0 flex-col border-r bg-surface lg:flex">
        <div className="flex h-10 shrink-0 items-center justify-between px-4 text-xs font-medium text-muted-foreground"><span>{doc.name}</span><span>{doc.pages.length} pages</span></div>
        <ol className="flex min-h-0 flex-1 flex-col gap-3 overflow-y-auto px-4 pb-4">
          {doc.pages.map((p, i) => {
            const n = i + 1;
            const active = activePage === n;
            return (
              <li key={n} id={`page-${n}`} ref={(el) => { pageRefs.current[n] = el; }}>
                <button type="button" onClick={() => setActivePage(n)} className={cn("flex aspect-[3/4] w-full flex-col rounded-lg border bg-background p-3 text-left shadow-xs transition-all", active ? "border-brand ring-2 ring-brand/30" : "hover:border-brand/40")}>
                  <span className={cn("mb-2 line-clamp-6 text-[9px] leading-[13px]", active ? "text-foreground" : "text-muted-foreground/70")}>
                    {active ? <mark className="rounded bg-brand-soft px-0.5 text-foreground">{p}</mark> : p}
                  </span>
                  <span className="mt-auto space-y-1">{[90, 100, 70, 95, 60].map((w, j) => <span key={j} className="block h-1 rounded bg-muted" style={{ width: `${w}%` }} />)}</span>
                  <span className="mt-2 text-[10px] tabular-nums text-muted-foreground">{n}</span>
                </button>
              </li>
            );
          })}
        </ol>
      </section>

      <section className="flex min-h-0 flex-col">
        <header className="flex h-12 shrink-0 items-center gap-2 border-b px-4 text-sm font-medium">
          <span className="flex size-6 items-center justify-center rounded-md bg-brand text-primary-foreground"><DocumentTextIcon size={14} className="flex" /></span>
          Ask your documents
          {activePage ? <span className="ml-auto text-xs font-normal text-muted-foreground">Viewing page {activePage}</span> : null}
        </header>
        <div className="relative min-h-0 flex-1" onClickCapture={onCitationClick}>
          {turns.length === 0 && !pending ? (
            <div className="flex h-full flex-col items-center justify-center gap-5 px-6 text-center">
              <p className="max-w-sm text-sm text-muted-foreground">Answers cite the page they came from. Hover a citation to preview it, click to jump to the page.</p>
              <Suggestions onSelect={ask}><SuggestionList className="justify-center">{suggestions.map((s) => <Suggestion key={s}>{s}</Suggestion>)}</SuggestionList></Suggestions>
            </div>
          ) : (
            <Thread className="h-full">
              <ThreadContent className="gap-5 p-4">
                {turns.map((t) =>
                  t.role === "user" ? (
                    <Message key={t.id} from="user"><MessageStack><MessageContent>{t.text}</MessageContent></MessageStack></Message>
                  ) : (
                    <Message key={t.id} from="assistant" className="max-w-none">
                      <MessageAvatar fallback={<SparklesIcon size={14} className="flex" />} />
                      <MessageStack><MessageContent><CitedText text={t.text} sources={sourcesFor(t.pages ?? [])} activeId={activeSourceId} onHover={(id) => { const m = /-p(\d+)$/.exec(id ?? ""); if (m) setActivePage(Number(m[1])); }} /></MessageContent></MessageStack>
                    </Message>
                  ),
                )}
                {pending ? (
                  <Message from="assistant" className="max-w-none">
                    <MessageAvatar fallback={<SparklesIcon size={14} className="flex" />} />
                    <MessageStack><MessageContent>{text ? <CitedText text={text} sources={sourcesFor(pending)} activeId={activeSourceId} /> : <TextShimmer className="text-sm text-muted-foreground" invertLight>Reading {doc.name}…</TextShimmer>}</MessageContent></MessageStack>
                  </Message>
                ) : null}
              </ThreadContent>
              <ThreadScrollToBottom />
            </Thread>
          )}
        </div>
        <div className="shrink-0 p-3 pt-0">
          <PromptInput onSubmit={ask}>
            <PromptInputTextarea value={draft} onChange={(e) => setDraft(e.target.value)} placeholder={`Ask about ${doc.name}`} />
            <PromptInputActions>
              <PromptInputActionGroup><span className="px-2 text-xs text-muted-foreground">{files.length} document{files.length === 1 ? "" : "s"} indexed</span></PromptInputActionGroup>
              <PromptInputActionGroup>
                {isStreaming ? (
                  <PromptInputAction asChild tooltip="Stop"><Button size="icon-sm" variant="secondary" className="rounded-full" aria-label="Stop" onClick={() => { stop(); setPending(null); }}><StopIcon size={16} className="flex" /></Button></PromptInputAction>
                ) : (
                  <PromptInputAction asChild tooltip={{ content: "Ask", shortcut: "↵" }}><Button size="icon-sm" className="rounded-full" aria-label="Ask" disabled={!draft.trim()} onClick={() => ask(draft)}><ArrowUpIcon size={16} className="flex" /></Button></PromptInputAction>
                )}
              </PromptInputActionGroup>
            </PromptInputActions>
          </PromptInput>
        </div>
      </section>
    </div>
  );
}

/**
 * Helpers for scripted streaming in demos. Blocks use these to fake a model
 * response token by token without touching the network. Swap them for a real
 * fetch in production; the UI code does not change.
 */

export function sleep(ms: number, signal?: AbortSignal) {
  return new Promise<void>((resolve, reject) => {
    if (signal?.aborted) return reject(signal.reason);
    const id = setTimeout(resolve, ms);
    signal?.addEventListener(
      "abort",
      () => {
        clearTimeout(id);
        reject(signal.reason);
      },
      { once: true },
    );
  });
}

/** Split text into word-ish chunks so streaming looks natural. */
export function chunk(text: string): string[] {
  return text.match(/\S+\s*|\s+/g) ?? [text];
}

export function jitter(base: number) {
  return base + Math.random() * base;
}

/** Yield a string piece by piece with a small random delay between pieces. */
export async function* streamText(
  text: string,
  signal?: AbortSignal,
  delayMs = 14,
): AsyncGenerator<string> {
  for (const piece of chunk(text)) {
    if (signal?.aborted) return;
    yield piece;
    await sleep(jitter(delayMs), signal);
  }
}

let counter = 0;
export const uid = (prefix = "id") => `${prefix}-${Date.now().toString(36)}-${++counter}`;

"use client";

import * as React from "react";

import { streamText } from "./mock-stream";

/**
 * Drive a piece of text onto the screen as if it were streaming from a model.
 * `start(text)` clears and streams; `stop()` freezes what has arrived.
 */
export function useStreamText(delayMs = 14) {
  const [text, setText] = React.useState("");
  const [isStreaming, setIsStreaming] = React.useState(false);
  const abortRef = React.useRef<AbortController | null>(null);

  const stop = React.useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
    setIsStreaming(false);
  }, []);

  const start = React.useCallback(
    async (full: string, onDone?: () => void) => {
      abortRef.current?.abort();
      const controller = new AbortController();
      abortRef.current = controller;
      setText("");
      setIsStreaming(true);
      try {
        for await (const piece of streamText(full, controller.signal, delayMs)) {
          setText((prev) => prev + piece);
        }
        if (!controller.signal.aborted) onDone?.();
      } catch {
        /* aborted */
      } finally {
        if (abortRef.current === controller) {
          abortRef.current = null;
          setIsStreaming(false);
        }
      }
    },
    [delayMs],
  );

  const reset = React.useCallback(() => {
    stop();
    setText("");
  }, [stop]);

  React.useEffect(() => () => abortRef.current?.abort(), []);

  return { text, isStreaming, start, stop, reset, setText };
}