frontend/src/components/EmbedStudio.tsx

import { useState, useEffect, useRef, useMemo, useCallback, type KeyboardEvent } from "react";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";
import { X, Send, Square, Code, Eye, Plus, Loader } from "lucide-react";
import type { UIMessage } from "ai";
import type { EmbedStore } from "../editor/embeds/embed-store";
import type { EmbedDataStore } from "../editor/embeds/embed-data-store";
import { buildDoc } from "../editor/embeds/build-doc";
import { useEmbedChat } from "../hooks/useEmbedChat";
import { useEmbedData } from "../hooks/useEmbedData";
import { useTheme } from "../hooks/useTheme";
import { isToolPart, normalizeToolPart } from "../utils/ai-tool-parts";
import { FilesSidebar } from "./FilesSidebar";
import { DataFileViewer } from "./DataFileViewer";

interface EmbedStudioProps {
  src: string;
  embedStore: EmbedStore | null;
  dataStore?: EmbedDataStore | null;
  modelRef: React.RefObject<string>;
  userId: string;
  onClose: () => void;
  /**
   * Invoked when the agent chooses a descriptive filename for a brand
   * new chart. Parent should update the htmlEmbed node in the doc and
   * lift the new src into its own state so subsequent actions (close,
   * reopen, etc.) reference the renamed file.
   */
  onRename?: (oldSrc: string, newSrc: string) => void;
  /**
   * Invoked when the user picks a different chart from the sidebar.
   * Parent is responsible for updating the studio's `src` and (if it
   * wants a fresh chat scope) rotating the React key on this
   * component so persisted messages reload for the new file.
   */
  onSelectChart?: (name: string) => void;
}

const EMBED_TOOL_LABELS: Record<string, [string, string]> = {
  createEmbed: ["Creating chart...", "Created chart"],
  patchEmbed: ["Updating chart...", "Updated chart"],
  readEmbed: ["Reading chart...", "Read chart"],
  listDataFiles: ["Listing data files...", "Listed data files"],
  readDataFile: ["Reading data file...", "Read data file"],
};

function embedToolLabel(name: string, state: string): string {
  const pair = EMBED_TOOL_LABELS[name];
  if (!pair) return name;
  return state === "result" ? pair[1] : pair[0];
}

function toolSubtitle(
  name: string,
  input: unknown,
): string | null {
  if (!input || typeof input !== "object") return null;
  const args = input as Record<string, unknown>;
  switch (name) {
    case "createEmbed":
      return typeof args.title === "string" ? String(args.title) : null;
    case "readDataFile":
      return typeof args.name === "string" ? String(args.name) : null;
    case "patchEmbed": {
      const search = typeof args.search === "string" ? args.search : "";
      const firstLine = search.split("\n")[0]?.trim() ?? "";
      return firstLine ? firstLine.slice(0, 60) : null;
    }
    default:
      return null;
  }
}

function MessageBubble({ message }: { message: UIMessage }) {
  const isUser = message.role === "user";
  return (
    <div className={`es-message ${isUser ? "es-message--user" : "es-message--ai"}`}>
      {message.parts.map((part, i) => {
        if (part.type === "text") {
          return (
            <div key={i} className="es-message__text">
              {isUser ? part.text : (
                <ReactMarkdown remarkPlugins={[remarkGfm]}>{part.text}</ReactMarkdown>
              )}
            </div>
          );
        }
        if (isToolPart(part)) {
          const tool = normalizeToolPart(part);
          if (!tool) return null;
          const subtitle = toolSubtitle(tool.toolName, tool.input);
          const isDone = tool.state === "result";
          return (
            <div
              key={i}
              className={`es-message__tool ${isDone ? "es-message__tool--done" : "es-message__tool--running"}`}
            >
              {!isDone && <Loader size={11} className="spin" />}
              <span className="es-message__tool-name">
                {embedToolLabel(tool.toolName, tool.state)}
              </span>
              {subtitle && (
                <span className="es-message__tool-subtitle">{subtitle}</span>
              )}
            </div>
          );
        }
        return null;
      })}
    </div>
  );
}

function EmbedAgentStatus({ messages }: { messages: UIMessage[] }) {
  const last = messages[messages.length - 1];
  if (!last) return null;

  if (last.role === "assistant") {
    const parts = last.parts ?? [];
    const runningTools = parts
      .map((p) => normalizeToolPart(p))
      .filter((t): t is NonNullable<typeof t> => t !== null && t.state !== "result")
      .map((t) => t.toolName);

    if (runningTools.length > 0) {
      const pair = EMBED_TOOL_LABELS[runningTools[runningTools.length - 1]];
      const text = pair ? pair[0] : runningTools[runningTools.length - 1];
      return (
        <div className="chat-panel__thinking">
          <Loader size={13} className="spin" />
          <span className="shimmer-text">{text}</span>
        </div>
      );
    }

    const hasText = parts.some((p: any) => p.type === "text" && p.text?.trim());
    if (hasText) return null;
  }

  return (
    <div className="chat-panel__thinking">
      <Loader size={13} className="spin" />
      <span className="shimmer-text">Thinking...</span>
    </div>
  );
}

export function EmbedStudio({
  src,
  embedStore,
  dataStore = null,
  modelRef,
  userId,
  onClose,
  onRename,
  onSelectChart,
}: EmbedStudioProps) {
  const { isDark, primaryColor } = useTheme();
  // The banner is a special embed: it renders edge-to-edge in a 5:2 box (the
  // hero) rather than as a height-reporting chart. The studio must preview it
  // the same way (fullBleed srcdoc + constrained 5:2 stage) so it's WYSIWYG.
  const isBanner = src === "banner.html";
  const [html, setHtml] = useState("");
  const [viewMode, setViewMode] = useState<"preview" | "code">("preview");
  const [selectedDataFile, setSelectedDataFile] = useState<string | null>(null);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const textareaRef = useRef<HTMLTextAreaElement>(null);

  // Iframe double-buffer
  const frameA = useRef<HTMLIFrameElement>(null);
  const frameB = useRef<HTMLIFrameElement>(null);
  const [activeFrame, setActiveFrame] = useState<"a" | "b">("a");
  const activeRef = useRef<"a" | "b">("a");
  activeRef.current = activeFrame;

  const data = useEmbedData({ dataStore, userId });
  const chat = useEmbedChat({ embedStore, dataStore, src, modelRef, userId, isDark, onRename });

  const selectedFile = selectedDataFile ? data.getFile(selectedDataFile) : undefined;

  // Sync from embed store
  useEffect(() => {
    if (!embedStore || !src) return;
    setHtml(embedStore.get(src));
    return embedStore.observeKey(src, setHtml);
  }, [embedStore, src]);

  // Build srcdoc. Banner previews use fullBleed so the content fills the
  // iframe edge-to-edge, matching FrontmatterHero and the published output.
  const srcdoc = useMemo(() => {
    if (!html) return "";
    return buildDoc(html, { isDark, primaryColor, fullBleed: isBanner });
  }, [html, isDark, primaryColor, isBanner]);

  // Load srcdoc into inactive frame
  useEffect(() => {
    if (!srcdoc) return;
    const next = activeRef.current === "a" ? "b" : "a";
    const frame = next === "a" ? frameA.current : frameB.current;
    if (frame) frame.srcdoc = srcdoc;
  }, [srcdoc]);

  const handleFrameLoad = useCallback((slot: "a" | "b") => {
    if (slot !== activeRef.current) {
      setActiveFrame(slot);
    }
  }, []);

  // Auto-scroll messages
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [chat.messages]);

  // Dev-only hook: let demo scripts inject a fake assistant reply into
  // the embed chat so the recording shows the agent "responding"
  // without waiting for a real LLM round-trip.
  useEffect(() => {
    if (!import.meta.env.DEV) return;
    const handler = (event: Event) => {
      const detail = (event as CustomEvent).detail as
        | { messages?: UIMessage[] }
        | undefined;
      if (Array.isArray(detail?.messages)) {
        chat.setMessages(detail.messages);
      }
    };
    window.addEventListener("__demo-embed-chat", handler);
    return () => window.removeEventListener("__demo-embed-chat", handler);
  }, [chat]);

  const handleSend = useCallback(() => {
    const text = chat.input.trim();
    if (!text || chat.isLoading) return;
    chat.sendMessage(text);
    chat.setInput("");
  }, [chat]);

  const handleKeyDown = useCallback(
    (e: KeyboardEvent<HTMLTextAreaElement>) => {
      if (e.key === "Enter" && !e.shiftKey) {
        e.preventDefault();
        handleSend();
      }
    },
    [handleSend],
  );

  const handleTextareaInput = useCallback(() => {
    const ta = textareaRef.current;
    if (ta) {
      ta.style.height = "auto";
      ta.style.height = Math.min(ta.scrollHeight, 160) + "px";
    }
  }, []);

  const frameStyle = (slot: "a" | "b"): React.CSSProperties => ({
    position: "absolute",
    inset: 0,
    border: "none",
    width: "100%",
    height: "100%",
    transition: "opacity 120ms ease",
    opacity: activeFrame === slot ? 1 : 0,
    zIndex: activeFrame === slot ? 2 : 1,
  });

  return (
    <div className="embed-studio">
      {/* Header */}
      <div className="es-header">
        <div className="es-header__left">
          <span className="es-header__icon">📊</span>
          <span className="es-header__title">Embed Studio</span>
          <code className="es-header__src">{src}</code>
        </div>
        <div className="es-header__actions">
          <button
            className="embed-btn"
            onClick={() => chat.clearMessages()}
            aria-label="New conversation"
            title="New conversation"
          >
            <Plus size={16} />
          </button>
          <button
            className="embed-btn"
            onClick={onClose}
            aria-label="Close Embed Studio"
          >
            <X size={16} />
          </button>
        </div>
      </div>

      {/* Body: split panel */}
      <div className="es-body">
        {/* Far-left: Files */}
        <FilesSidebar
          embedStore={embedStore}
          currentSrc={src}
          dataFiles={data.files}
          selectedDataFile={selectedDataFile}
          onSelectChart={onSelectChart}
          onSelectDataFile={setSelectedDataFile}
          onUploadFiles={data.uploadFiles}
          onRemoveDataFile={data.removeFile}
        />

        {/* Left: Chat */}
        <div className="es-chat">
          <div className="es-chat__messages">
            {chat.messages.length === 0 && (() => {
              const hasChart = html.trim().length > 0;
              let title: string;
              let hint: string;
              if (isBanner) {
                title = hasChart
                  ? "Ask for changes to the banner."
                  : "Describe the banner you want.";
                hint = hasChart
                  ? 'Example: "make it slower and more abstract"'
                  : 'Example: "abstract flow field in the primary color, very subtle motion"';
              } else if (hasChart) {
                title = "Ask for changes to the chart.";
                hint =
                  'Example: "switch to a sequential palette" or "highlight the biggest bar"';
              } else {
                title = "Describe the chart you want to create.";
                hint =
                  'Example: "Create a bar chart showing model sizes with categorical colors"';
              }
              return (
                <div className="es-chat__empty">
                  <p>{title}</p>
                  <p className="es-chat__hint">{hint}</p>
                </div>
              );
            })()}
            {chat.messages.map((msg) => (
              <MessageBubble key={msg.id} message={msg} />
            ))}
            {chat.isLoading && chat.messages.length > 0 && (
              <EmbedAgentStatus messages={chat.messages} />
            )}
            <div ref={messagesEndRef} />
          </div>

          {/* Input */}
          <div className="es-chat__input-area">
            {chat.error && (
              <div className="es-chat__error">
                {chat.error.message || "An error occurred"}
              </div>
            )}
            <div className="es-chat__input-row">
              <textarea
                ref={textareaRef}
                value={chat.input}
                onChange={(e) => chat.setInput(e.target.value)}
                onKeyDown={handleKeyDown}
                onInput={handleTextareaInput}
                placeholder={
                  html.trim().length > 0
                    ? src === "banner.html"
                      ? "Tweak the banner..."
                      : "Tweak the chart..."
                    : src === "banner.html"
                      ? "Describe your banner..."
                      : "Describe your chart..."
                }
                rows={1}
                className="es-chat__textarea"
              />
              {chat.isLoading ? (
                <button
                  className="es-chat__send-btn"
                  onClick={chat.stop}
                  aria-label="Stop"
                >
                  <Square size={16} />
                </button>
              ) : (
                <button
                  className="es-chat__send-btn"
                  onClick={handleSend}
                  disabled={!chat.input.trim()}
                  aria-label="Send"
                >
                  <Send size={16} />
                </button>
              )}
            </div>
          </div>
        </div>

        {/* Right: Preview (chart) or data file viewer */}
        <div className="es-preview">
          {selectedFile ? (
            <DataFileViewer
              file={selectedFile}
              onClose={() => setSelectedDataFile(null)}
            />
          ) : null}

          <div
            className={`es-preview__toolbar ${selectedFile ? "es-hidden" : ""}`}
          >
            <button
              className={`es-preview__tab ${viewMode === "preview" ? "es-preview__tab--active" : ""}`}
              onClick={() => setViewMode("preview")}
            >
              <Eye size={14} />
              Preview
            </button>
            <button
              className={`es-preview__tab ${viewMode === "code" ? "es-preview__tab--active" : ""}`}
              onClick={() => setViewMode("code")}
            >
              <Code size={14} />
              Code
            </button>
            <div style={{ flex: 1 }} />
            <button className="embed-btn embed-btn-primary" onClick={onClose}>
              Save & Close
            </button>
          </div>

          {/* Both panels always mounted - toggle with hidden class to preserve iframe state */}
          <div
            className={`es-preview__frame-container ${isBanner ? "es-preview__frame-container--banner" : ""} ${viewMode !== "preview" || selectedFile ? "es-hidden" : ""}`}
          >
            {html ? (
              <div
                className={`es-preview__stage ${isBanner ? "es-preview__stage--banner" : ""}`}
              >
                <iframe
                  ref={frameA}
                  title="Chart preview A"
                  sandbox="allow-scripts allow-same-origin"
                  onLoad={() => handleFrameLoad("a")}
                  style={frameStyle("a")}
                />
                <iframe
                  ref={frameB}
                  title="Chart preview B"
                  sandbox="allow-scripts allow-same-origin"
                  onLoad={() => handleFrameLoad("b")}
                  style={frameStyle("b")}
                />
              </div>
            ) : (
              <div className="es-preview__empty">
                <span style={{ fontSize: 32, opacity: 0.4 }}>📊</span>
                <p>Chart preview will appear here</p>
              </div>
            )}
          </div>
          <div
            className={`es-preview__code ${viewMode !== "code" || selectedFile ? "es-hidden" : ""}`}
          >
            <pre className="es-preview__pre">
              <code>{html || "// No chart HTML yet"}</code>
            </pre>
          </div>
        </div>
      </div>
    </div>
  );
}