AI Chat

"use client";

import { useRef, useState } from "react";
import { CitationStack, type Citation } from "@/lib/citations";
import {
  ExaContentAccumulator,
  streamNdjson,
  type ExaStreamEvent,
} from "@/lib/exa-stream";

type UserMessage = { id: string; role: "user"; content: string };

type AssistantMessage = {
  id: string;
  role: "assistant";
  content: string;
  status: "streaming" | "done" | "error";
  citations: Citation[];
  error?: string;
};

type Message = UserMessage | AssistantMessage;

const SUGGESTIONS = [
  "Latest news on Nvidia",
  "What is Stripe working on right now?",
  "Recent buying signals at Anthropic",
];

export function ChatExample() {
  const [messages, setMessages] = useState<Message[]>([]);
  const [draft, setDraft] = useState("");
  const [streaming, setStreaming] = useState(false);
  const idRef = useRef(0);

  function nextId(prefix: string) {
    idRef.current += 1;
    return `${prefix}_${idRef.current}`;
  }

  function patchAssistant(id: string, patch: Partial<AssistantMessage>) {
    setMessages((prev) =>
      prev.map((m) =>
        m.id === id && m.role === "assistant" ? { ...m, ...patch } : m
      )
    );
  }

  async function send(text: string) {
    const query = text.trim();
    if (!query || streaming) return;
    setDraft("");
    const userMsg: UserMessage = { id: nextId("u"), role: "user", content: query };
    const asstId = nextId("a");
    const asstMsg: AssistantMessage = {
      id: asstId,
      role: "assistant",
      content: "",
      status: "streaming",
      citations: [],
    };
    const history = messages
      .filter((m) =>
        m.role === "user" ? true : m.status === "done" && m.content
      )
      .slice(-6)
      .map((m) => ({ role: m.role, content: m.content }));
    setMessages((m) => [...m, userMsg, asstMsg]);
    setStreaming(true);

    try {
      const res = await fetch("/api/chat", {
        method: "POST",
        headers: { "content-type": "application/json" },
        body: JSON.stringify({ query, history }),
      });
      if (!res.ok || !res.body) {
        const text = await res.text().catch(() => "");
        let err = `Request failed (${res.status})`;
        try {
          err = JSON.parse(text).error ?? err;
        } catch {}
        patchAssistant(asstId, { status: "error", error: err });
        return;
      }

      const acc = new ExaContentAccumulator();
      const citationsById = new Map<string, Citation>();

      await streamNdjson<ExaStreamEvent>(res, (event) => {
        if (event.type === "content") {
          patchAssistant(asstId, { content: acc.push(event.text) });
        } else if (event.type === "citations") {
          for (const c of event.citations) citationsById.set(c.id, c);
          patchAssistant(asstId, {
            citations: Array.from(citationsById.values()),
          });
        } else if (event.type === "done") {
          patchAssistant(asstId, { status: "done" });
        } else if (event.type === "error") {
          patchAssistant(asstId, { status: "error", error: event.error });
        }
      });
    } catch (err) {
      patchAssistant(asstId, {
        status: "error",
        error: err instanceof Error ? err.message : String(err),
      });
    } finally {
      setStreaming(false);
    }
  }

  return (
    <div className="border border-[#e5e5e5] bg-white overflow-hidden flex flex-col">
      <div className="px-4 py-4 space-y-4 min-h-[280px] max-h-[520px] overflow-y-auto">
        {messages.length === 0 ? (
          <EmptyState onPick={(q) => void send(q)} />
        ) : (
          messages.map((m) =>
            m.role === "user" ? (
              <UserBubble key={m.id} message={m} />
            ) : (
              <AssistantBubble key={m.id} message={m} />
            )
          )
        )}
      </div>
      <form
        onSubmit={(e) => {
          e.preventDefault();
          void send(draft);
        }}
        className="flex items-center gap-2 border-t border-[#e5e5e5] bg-[#fafafa] px-3 py-2"
      >
        <input
          type="text"
          value={draft}
          onChange={(e) => setDraft(e.target.value)}
          placeholder="Ask something about a company, person, or market..."
          className="flex-1 border border-[#e5e5e5] bg-white text-[13px] text-[#111827] px-3 py-2 focus:outline-none focus:border-[var(--color-brand)]"
        />
        <button
          type="submit"
          disabled={streaming || !draft.trim()}
          className="bg-[var(--color-brand)] text-white text-[13px] font-medium px-3.5 py-2 disabled:opacity-50"
        >
          {streaming ? "Thinking..." : "Send"}
        </button>
      </form>
    </div>
  );
}

function EmptyState({ onPick }: { onPick: (q: string) => void }) {
  return (
    <div className="text-[13px] text-[#6b7280]">
      <p className="mb-3">Try one of these:</p>
      <div className="flex flex-wrap gap-2">
        {SUGGESTIONS.map((s) => (
          <button
            key={s}
            type="button"
            onClick={() => onPick(s)}
            className="rounded-full border border-[#e5e5e5] bg-white px-3 py-1 text-[12px] text-[#374151] hover:border-[var(--color-brand)] hover:text-[#111827]"
          >
            {s}
          </button>
        ))}
      </div>
    </div>
  );
}

function UserBubble({ message }: { message: UserMessage }) {
  return (
    <div className="flex justify-end">
      <div className="max-w-[80%] bg-[var(--color-brand)] text-white text-[13px] px-3 py-2 leading-[1.55]">
        {message.content}
      </div>
    </div>
  );
}

function AssistantBubble({ message }: { message: AssistantMessage }) {
  const paragraphs = message.content ? message.content.split(/\n{2,}/) : [];
  return (
    <div className="space-y-2">
      {paragraphs.length === 0 && message.status === "streaming" && (
        <div className="text-[13px]">
          <span className="inline-block h-3 w-24 rounded bg-[#e5e7eb] animate-pulse align-middle" />
        </div>
      )}
      {paragraphs.map((p, i) => (
        <div
          key={i}
          className="text-[13px] text-[#111827] leading-[1.65] whitespace-pre-line"
        >
          <RenderContent
            content={p}
            citations={message.citations}
            streaming={
              message.status === "streaming" && i === paragraphs.length - 1
            }
          />
        </div>
      ))}
      {message.status === "error" && (
        <div className="text-[12px] text-[#b91c1c]">{message.error}</div>
      )}
    </div>
  );
}

type Segment =
  | { kind: "text"; text: string }
  | { kind: "citations"; indices: number[] };

function parseSegments(content: string): Segment[] {
  const out: Segment[] = [];
  const re = /(?:\[(\d+)\])+/g;
  let last = 0;
  let m: RegExpExecArray | null;
  while ((m = re.exec(content)) !== null) {
    if (m.index > last)
      out.push({ kind: "text", text: content.slice(last, m.index) });
    const indices = Array.from(m[0].matchAll(/\[(\d+)\]/g)).map((x) =>
      parseInt(x[1], 10)
    );
    out.push({ kind: "citations", indices });
    last = re.lastIndex;
  }
  if (last < content.length)
    out.push({ kind: "text", text: content.slice(last) });
  return out;
}

function RenderContent({
  content,
  citations,
  streaming,
}: {
  content: string;
  citations: Citation[];
  streaming: boolean;
}) {
  return (
    <>
      {parseSegments(content).map((seg, i) =>
        seg.kind === "text" ? (
          <span key={i}>{seg.text}</span>
        ) : (
          <CitationStack key={i} indices={seg.indices} citations={citations} />
        )
      )}
      {streaming && (
        <span className="inline-block h-3 w-1.5 ml-0.5 bg-[#9ca3af] animate-pulse align-middle" />
      )}
    </>
  );
}