frontend/src/hooks/useEmbedChat.ts

import { useChat } from "@ai-sdk/react";
import {
  DefaultChatTransport,
  lastAssistantMessageIsCompleteWithToolCalls,
} from "ai";
import { useCallback, useEffect, useRef, useState } from "react";
import type { UIMessage } from "ai";
import type { EmbedStore } from "../editor/embeds/embed-store";
import type { EmbedDataStore } from "../editor/embeds/embed-data-store";
import { loadMessages, saveMessages } from "../utils/chat-persistence";

interface UseEmbedChatOptions {
  embedStore: EmbedStore | null;
  dataStore?: EmbedDataStore | null;
  /** Current embed src key (filename) */
  src: string;
  modelRef: React.RefObject<string>;
  userId: string;
  /**
   * Current article theme (light/dark). Forwarded to the agent so it
   * can sanity-check its colour choices and bake theme-aware CSS into
   * generated charts (instead of hardcoding axis/text colours that
   * only happen to look right under one theme).
   */
  isDark?: boolean;
  /**
   * Called when the agent picks a descriptive filename via createEmbed.
   * Parent is expected to update the htmlEmbed node's src attribute in
   * the doc and lift the new src into its own state. The hook already
   * moves the content in embedStore and migrates persisted messages.
   */
  onRename?: (oldSrc: string, newSrc: string) => void;
}

/**
 * Slugify an arbitrary filename hint into a safe kebab-case slug.
 * Keeps ASCII letters/digits/dashes, collapses whitespace, lowercases.
 */
function slugify(raw: string): string {
  const base = raw
    .toLowerCase()
    .replace(/\.html?$/i, "")
    .normalize("NFKD")
    .replace(/[\u0300-\u036f]/g, "")
    .replace(/[^a-z0-9]+/g, "-")
    .replace(/^-+|-+$/g, "")
    .slice(0, 60);
  return base || "chart";
}

/**
 * Given a slug, find a unique .html filename in `embedStore`. If the
 * desired name is free, return it; otherwise append `-2`, `-3`, ...
 */
function ensureUniqueFilename(
  slug: string,
  embedStore: EmbedStore,
  currentSrc: string,
): string {
  const base = `${slug}.html`;
  if (base === currentSrc) return base;
  if (!embedStore.has(base)) return base;
  let i = 2;
  while (embedStore.has(`${slug}-${i}.html`)) i++;
  return `${slug}-${i}.html`;
}

const MAX_EMBED_SIZE = 512_000;
/**
 * Hard cap on the number of characters we inline when the agent reads
 * a data file. Above this, the tool returns a truncated preview with
 * a warning so the model doesn't blow past its context window on a
 * single tool call.
 */
const MAX_DATA_INLINE = 120_000;

const transport = new DefaultChatTransport({ api: "/api/embed-chat" });

function scopeKey(src: string): string {
  return `embed:${src}`;
}

export function useEmbedChat({ embedStore, dataStore, src, modelRef, userId, isDark, onRename }: UseEmbedChatOptions) {
  const [input, setInput] = useState("");
  // srcRef is the source of truth for the active embed filename. It
  // starts from the `src` prop but may change mid-session when the
  // agent renames the chart via createEmbed({ filename }). We do NOT
  // overwrite it from the prop on every render to avoid clobbering a
  // freshly-chosen name before the parent has caught up.
  const srcRef = useRef(src);
  // Sync srcRef whenever the incoming `src` prop changes from outside.
  // After an internal rename, we already set srcRef to the new name and
  // asked the parent to lift it; when the parent's state arrives here
  // the values match and this is a no-op.
  if (src !== srcRef.current) {
    srcRef.current = src;
  }
  const userIdRef = useRef(userId);
  userIdRef.current = userId;
  const onRenameRef = useRef(onRename);
  onRenameRef.current = onRename;

  // Guard: only allow sendAutomaticallyWhen AFTER the user has sent at
  // least one message in this session. Prevents auto-send on mount when
  // restored messages end with completed tool calls.
  const userHasSentRef = useRef(false);

  // Track the current theme via ref so getEmbedContext always reads the
  // latest value without forcing a callback identity churn on toggle.
  const isDarkRef = useRef(!!isDark);
  isDarkRef.current = !!isDark;

  const getEmbedContext = useCallback(() => {
    if (!embedStore || !srcRef.current) return {};
    const currentSrc = srcRef.current;
    const embedHtml = embedStore.get(currentSrc);
    const isBanner = currentSrc === "banner.html";
    // A "scratch" file is a newly-created placeholder whose name is a
    // hash-based auto-generated slug (see SlashMenu). When the chart is
    // still empty under such a name, we explicitly ask the agent to
    // pick a descriptive filename on the first createEmbed call.
    const isScratch =
      !isBanner && !embedHtml && /^d3-chart-[a-z0-9]+\.html$/i.test(currentSrc);
    const dataFiles = dataStore
      ? dataStore.list().map((m) => ({
          name: m.name,
          ext: m.ext,
          size: m.size,
          rowCount: m.rowCount,
          columns: m.columns,
        }))
      : [];
    return {
      embedHtml: embedHtml || undefined,
      isBanner,
      isScratch,
      dataFiles,
      theme: isDarkRef.current ? "dark" : "light",
    };
  }, [embedStore, dataStore]);

  const executeToolCall = useCallback(
    (toolCall: { toolName: string; args: unknown; toolCallId: string }) => {
      if (!embedStore) return "Embed store not available";
      const currentSrc = srcRef.current;
      if (!currentSrc) return "No embed src specified";

      switch (toolCall.toolName) {
        case "createEmbed": {
          const { html, title, filename } = toolCall.args as {
            html: string;
            title?: string;
            source?: string;
            filename?: string;
          };
          if (!html?.trim()) return "ERROR: html cannot be empty";
          const stripped = html.trim();
          if (stripped.length > MAX_EMBED_SIZE) {
            return `ERROR: chart too large (${stripped.length} bytes, max ${MAX_EMBED_SIZE}). Simplify the code.`;
          }

          // Resolve the target src: if the agent proposed a filename
          // slug, slugify + ensure uniqueness. Never rename the banner
          // (its name is load-bearing for the article layout).
          let targetSrc = currentSrc;
          const isBanner = currentSrc === "banner.html";
          if (filename && !isBanner) {
            const slug = slugify(filename);
            const candidate = ensureUniqueFilename(slug, embedStore, currentSrc);
            if (candidate !== currentSrc) {
              targetSrc = candidate;
            }
          }

          if (targetSrc !== currentSrc) {
            // Write html under the new key and remove the old one atomically.
            embedStore.set(targetSrc, stripped);
            embedStore.remove(currentSrc);
            // Migrate persisted chat messages so the localStorage scope
            // tracks the new file name.
            try {
              const uid = userIdRef.current;
              const existing = loadMessages(uid, scopeKey(currentSrc));
              if (existing && existing.length > 0) {
                saveMessages(uid, scopeKey(targetSrc), existing);
              }
              saveMessages(uid, scopeKey(currentSrc), []);
            } catch {
              // Non-fatal: persistence is best-effort.
            }
            srcRef.current = targetSrc;
            // Let the parent update the htmlEmbed node's src attribute
            // in the ProseMirror doc and lift the new src into state.
            onRenameRef.current?.(currentSrc, targetSrc);
          } else {
            embedStore.set(currentSrc, stripped);
          }

          const lines = stripped.split("\n").length;
          const renamed = targetSrc !== currentSrc ? ` [renamed to ${targetSrc}]` : "";
          return `chart created (${lines} lines, ${stripped.length} chars)${title ? ` - "${title}"` : ""}${renamed}`;
        }

        case "patchEmbed": {
          const { search, replace } = toolCall.args as {
            search: string;
            replace: string;
          };
          const current = embedStore.get(currentSrc);
          if (!current) {
            return "ERROR: no chart exists yet. Use createEmbed first.";
          }
          if (!search) {
            return "ERROR: search block cannot be empty";
          }
          if (!current.includes(search)) {
            return (
              "ERROR: search block not found in current chart. " +
              "Call readEmbed to see the exact current content, then retry."
            );
          }
          const patched = current.replace(search, replace);
          if (patched.length > MAX_EMBED_SIZE) {
            return `ERROR: patch would exceed size limit (${patched.length} bytes, max ${MAX_EMBED_SIZE}).`;
          }
          embedStore.set(currentSrc, patched);
          return "patch applied successfully";
        }

        case "readEmbed": {
          const current = embedStore.get(currentSrc);
          if (!current) return "(no chart yet)";
          return current;
        }

        case "listDataFiles": {
          if (!dataStore) return "(no data files - data store unavailable)";
          const files = dataStore.list();
          if (files.length === 0) return "(no data files attached)";
          return files
            .map((f) => {
              const cols = f.columns && f.columns.length > 0
                ? ` columns=[${f.columns.join(", ")}]`
                : "";
              const rows = f.rowCount !== undefined ? ` rows=${f.rowCount}` : "";
              return `- ${f.name} (${f.ext}, ${f.size} bytes)${rows}${cols}`;
            })
            .join("\n");
        }

        case "readDataFile": {
          if (!dataStore) return "ERROR: data store unavailable";
          const { name } = toolCall.args as { name: string };
          if (!name) return "ERROR: missing file name";
          const file = dataStore.get(name);
          if (!file) return `ERROR: data file "${name}" not found`;
          const content = file.content;
          if (content.length <= MAX_DATA_INLINE) return content;
          return (
            content.slice(0, MAX_DATA_INLINE) +
            `\n\n[TRUNCATED - full file is ${content.length} chars, showing first ${MAX_DATA_INLINE}]`
          );
        }

        default:
          return `Unknown tool: ${toolCall.toolName}`;
      }
    },
    [embedStore, dataStore],
  );

  const { addToolOutput, setMessages: chatSetMessages, ...chat } = useChat({
    transport,
    // No initialMessages - we restore via setMessages below to match
    // the proven pattern used by the global chat.
    sendAutomaticallyWhen(ctx) {
      if (!userHasSentRef.current) return false;
      return lastAssistantMessageIsCompleteWithToolCalls(ctx);
    },

    async onToolCall({ toolCall }) {
      if (toolCall.dynamic) return;

      const result = executeToolCall({
        toolName: toolCall.toolName as string,
        args: toolCall.input,
        toolCallId: toolCall.toolCallId,
      });

      (addToolOutput as (args: { tool: string; toolCallId: string; output: unknown }) => void)({
        tool: toolCall.toolName as string,
        toolCallId: toolCall.toolCallId,
        output: result,
      });
    },

    onError(error) {
      console.error("[embed-chat] error:", error);
    },
  });

  // --- Restore persisted messages on mount (same pattern as global chat) ---
  const chatSetMessagesRef = useRef(chatSetMessages);
  chatSetMessagesRef.current = chatSetMessages;

  useEffect(() => {
    const key = scopeKey(src);
    const stored = loadMessages(userId, key);
    if (stored && stored.length > 0) {
      chatSetMessagesRef.current(stored);
    }
    // Only run once on mount (src/userId are stable per keyed instance)
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  // --- Persistence: save to localStorage ---------------------------------
  const messagesRef = useRef(chat.messages);
  messagesRef.current = chat.messages;

  const persist = useCallback(() => {
    const msgs = messagesRef.current as UIMessage[];
    const key = scopeKey(srcRef.current);
    const uid = userIdRef.current;
    if (msgs.length > 0) {
      saveMessages(uid, key, msgs);
    } else {
      saveMessages(uid, key, []);
    }
  }, []);

  // Save after each completed round (streaming/submitted -> ready)
  const prevStatusRef = useRef(chat.status);
  useEffect(() => {
    const prev = prevStatusRef.current;
    const curr = chat.status;
    prevStatusRef.current = curr;

    const wasActive = prev === "streaming" || prev === "submitted";
    if (wasActive && curr === "ready" && chat.messages.length > 0) {
      persist();
    }
  }, [chat.status, chat.messages, persist]);

  // Save on unmount (covers close-while-streaming)
  useEffect(() => () => persist(), [persist]);

  const sendMessage = useCallback(
    (content: string) => {
      userHasSentRef.current = true;
      const context = getEmbedContext();
      const model = modelRef.current;
      chat.sendMessage({ text: content }, { body: { context, model } });
    },
    [chat, getEmbedContext, modelRef],
  );

  const isLoading =
    chat.status === "streaming" || chat.status === "submitted";

  const clearMessages = useCallback(() => {
    chatSetMessages([]);
    saveMessages(userIdRef.current, scopeKey(srcRef.current), []);
  }, [chatSetMessages]);

  return {
    messages: chat.messages as UIMessage[],
    isLoading,
    error: chat.error,
    sendMessage,
    clearMessages,
    input,
    setInput,
    stop: chat.stop,
    /**
     * Imperative setter on the underlying chat. Exposed so dev-only demo
     * scripts can inject a fake assistant reply (used to record the demo
     * video without waiting for a real LLM round-trip). DO NOT use in
     * production code paths.
     */
    setMessages: chatSetMessages,
  };
}