astro-tractstack 2.0.14 → 2.0.16

package/templates/src/components/edit/pane/AddPanePanel_new.tsx

@@ -1,34 +1,110 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import { useEffect, useState, useMemo } from 'react';
-import { Switch } from '@ark-ui/react';
-import { Select } from '@ark-ui/react/select';
-import { Portal } from '@ark-ui/react/portal';
-import { createListCollection } from '@ark-ui/react/collection';
-import ChevronUpDownIcon from '@heroicons/react/20/solid/ChevronUpDownIcon';
-import CheckIcon from '@heroicons/react/20/solid/CheckIcon';
-import { NodesContext } from '@/stores/nodes';
-import {
-  PanesPreviewGenerator,
-  type PanePreviewRequest,
-  type PaneFragmentResult,
-} from '@/components/compositor/preview/PanesPreviewGenerator';
-import {
-  PaneSnapshotGenerator,
-  type SnapshotData,
-} from '@/components/compositor/preview/PaneSnapshotGenerator';
-import { createEmptyStorykeep } from '@/utils/compositor/nodesHelper';
+import { useState, useCallback } from 'react';
+import DocumentPlusIcon from '@heroicons/react/24/outline/DocumentPlusIcon';
+import SparklesIcon from '@heroicons/react/24/outline/SparklesIcon';
+import SwatchIcon from '@heroicons/react/24/outline/SwatchIcon';
+import { NodesContext, getCtx } from '@/stores/nodes';
 import { cloneDeep } from '@/utils/helpers';
-import {
-  brandColourStore,
-  preferredThemeStore,
-  hasAssemblyAIStore,
-} from '@/stores/storykeep';
-import { templateCategories } from '@/utils/compositor/templateMarkdownStyles';
-import { AiPaneGenerator } from './AiPaneGenerator';
-import { AddPaneNewCustomCopy } from './AddPanePanel_newCustomCopy';
-import { themes, type Theme, type BrandConfig } from '@/types/tractstack';
+import { hasAssemblyAIStore } from '@/stores/storykeep';
+import prompts from '@/constants/prompts.json';
+import type { BrandConfig, DesignLibraryEntry } from '@/types/tractstack';
 import { PaneAddMode, type TemplatePane } from '@/types/compositorTypes';
 import { useStore } from '@nanostores/react';
+import { CopyInputStep } from './steps/CopyInputStep';
+import { DesignLibraryStep } from './steps/DesignLibraryStep';
+import { AiDesignStep, type AiDesignConfig } from './steps/AiDesignStep';
+import { parseAiPane, parseAiCopyHtml } from '@/utils/compositor/aiPaneParser';
+import {
+  convertStorageToLiveTemplate,
+  mergeCopyIntoTemplate,
+  convertTemplateToAIShell,
+} from '@/utils/compositor/designLibraryHelper';
+
+type Step =
+  | 'initial'
+  | 'copyInput'
+  | 'designLibrary'
+  | 'aiDesign'
+  | 'loading'
+  | 'error';
+
+type InitialChoice = 'library' | 'ai' | 'blank';
+type CopyMode = 'prompt' | 'raw';
+
+interface GenerationResponse {
+  success: boolean;
+  data?: { response: string | object };
+  error?: string;
+}
+
+const callAskLemurAPI = async (
+  prompt: string,
+  context: string,
+  expectJson: boolean,
+  isSandboxMode: boolean
+): Promise<string> => {
+  const goBackend =
+    import.meta.env.PUBLIC_GO_BACKEND || 'http://localhost:8080';
+  const tenantId = import.meta.env.PUBLIC_TENANTID || 'default';
+  const requestBody = {
+    prompt,
+    input_text: context,
+    final_model: '',
+    temperature: 0.5,
+    max_tokens: 2000,
+  };
+
+  let response: Response;
+  if (isSandboxMode) {
+    response = await fetch(`/api/sandbox`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json', 'X-Tenant-ID': tenantId },
+      credentials: 'include',
+      body: JSON.stringify({ action: 'askLemur', payload: requestBody }),
+    });
+  } else {
+    response = await fetch(`${goBackend}/api/v1/aai/askLemur`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json', 'X-Tenant-ID': tenantId },
+      credentials: 'include',
+      body: JSON.stringify(requestBody),
+    });
+  }
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    let backendError = `API call failed: ${response.status} ${response.statusText}`;
+    try {
+      const errorJson = JSON.parse(errorText);
+      if (errorJson?.error) backendError = errorJson.error;
+    } catch (e) {
+      /* ignore */
+    }
+    throw new Error(backendError);
+  }
+
+  const result = (await response.json()) as GenerationResponse;
+  if (!result.success || !result.data?.response) {
+    throw new Error(
+      result.error || 'Generation failed to return valid response.'
+    );
+  }
+
+  let rawResponseData = result.data.response;
+  if (expectJson && typeof rawResponseData === 'object') {
+    return JSON.stringify(rawResponseData);
+  }
+  if (typeof rawResponseData === 'string') {
+    let responseString = rawResponseData;
+    if (responseString.startsWith('```json')) {
+      responseString = responseString.slice(7, -3).trim();
+    } else if (responseString.startsWith('```html')) {
+      responseString = responseString.slice(7, -3).trim();
+    }
+    return responseString;
+  }
+  throw new Error('Unexpected response format received from API.');
+};
 
 interface AddPaneNewPanelProps {
   nodeId: string;
@@ -38,244 +114,230 @@ interface AddPaneNewPanelProps {
   isStoryFragment?: boolean;
   isContextPane?: boolean;
   config?: BrandConfig;
+  isSandboxMode?: boolean;
 }
 
-interface PreviewPane {
-  ctx: NodesContext;
-  snapshot?: SnapshotData;
-  template: any;
-  index: number;
-  htmlFragment?: string;
-  fragmentError?: string;
-}
-
-interface TemplateCategory {
-  id: string;
-  title: string;
-  getTemplates: (theme: Theme, brand: string, useOdd: boolean) => any[];
-}
-
-const ITEMS_PER_PAGE = 8;
-
-type Mode = 'template' | 'custom' | 'ai';
-
 const AddPaneNewPanel = ({
   nodeId,
   first,
-  setMode: setParentMode, // Renamed prop to avoid conflict with local state
-  ctx,
+  setMode: setParentMode,
+  ctx: providedCtx,
   isStoryFragment = false,
   isContextPane = false,
   config,
+  isSandboxMode = false,
 }: AddPaneNewPanelProps) => {
-  const brand = useStore(brandColourStore);
+  const ctx = providedCtx || getCtx();
   const hasAssemblyAI = useStore(hasAssemblyAIStore);
-  const [mode, setMode] = useState<Mode>('template'); // Local mode state
-  const [customMarkdown, setCustomMarkdown] = useState<string>(`...`);
-  const [previews, setPreviews] = useState<PreviewPane[]>([]);
-  const [currentPage, setCurrentPage] = useState(0);
-  const [renderedPages, setRenderedPages] = useState<Set<number>>(new Set());
-  const [selectedTheme, setSelectedTheme] = useState<Theme>(
-    useStore(preferredThemeStore)
+  const [step, setStep] = useState<Step>('initial');
+  const [initialChoice, setInitialChoice] = useState<InitialChoice | null>(
+    null
   );
-  const [useOddVariant, setUseOddVariant] = useState(false);
-  const [selectedCategory, setSelectedCategory] = useState<TemplateCategory>(
-    templateCategories[isContextPane ? 1 : first ? 4 : 0]
-  );
-  const [isInserting, setIsInserting] = useState(false);
-  const [fragmentsToGenerate, setFragmentsToGenerate] = useState<
-    PanePreviewRequest[]
-  >([]);
-
-  const categoryCollection = useMemo(() => {
-    const categories = isContextPane
-      ? [templateCategories[1]]
-      : templateCategories;
-
-    return createListCollection({
-      items: categories,
-      itemToValue: (item) => item.id,
-      itemToString: (item) => item.title,
-    });
-  }, [isContextPane]);
-
-  const themesCollection = useMemo(() => {
-    return createListCollection({
-      items: themes,
-      itemToValue: (item) => item,
-      itemToString: (item) => item.replace(/-/g, ' '),
-    });
-  }, []);
-
-  const filteredTemplates = useMemo(() => {
-    return selectedCategory.getTemplates(selectedTheme, brand, useOddVariant);
-  }, [selectedTheme, useOddVariant, selectedCategory, brand]);
-
-  useEffect(() => {
-    if (isContextPane) {
-      setSelectedCategory(templateCategories[1]);
-    } else if (first) {
-      setSelectedCategory(templateCategories[4]);
+  const [error, setError] = useState<string | null>(null);
+  const [copyMode, setCopyMode] = useState<CopyMode>('prompt');
+  const [promptValue, setPromptValue] = useState('');
+  const [copyValue, setCopyValue] = useState('');
+  const [aiDesignConfig, setAiDesignConfig] = useState<AiDesignConfig>({
+    harmony: 'Analogous',
+    baseColor: '',
+    accentColor: '',
+    theme: 'Light',
+    additionalNotes: '',
+  });
+
+  const handleInitialChoice = (choice: InitialChoice) => {
+    setInitialChoice(choice);
+    setError(null);
+
+    if (choice === 'blank') {
+      handleBlankSlate();
     } else {
-      setSelectedCategory(templateCategories[0]);
+      setStep('copyInput');
     }
-  }, [isContextPane, first]);
-
-  useEffect(() => {
-    const newPreviews = filteredTemplates.map((template, index: number) => {
-      const previewCtx = new NodesContext();
-      previewCtx.addNode(createEmptyStorykeep('tmp'));
-      const thisTemplate =
-        mode === 'custom'
-          ? {
-              ...template,
-              markdown: template.markdown && {
-                ...template.markdown,
-                markdownBody: customMarkdown,
-              },
-            }
-          : template;
-      previewCtx.addTemplatePane('tmp', thisTemplate);
-      return { ctx: previewCtx, template: thisTemplate, index };
-    });
-    setPreviews(newPreviews);
-    setCurrentPage(0);
-    setRenderedPages(new Set());
-  }, [filteredTemplates, customMarkdown, mode]);
-
-  const totalPages = Math.ceil(previews.length / ITEMS_PER_PAGE);
-
-  const visiblePreviews = useMemo(() => {
-    const startIndex = currentPage * ITEMS_PER_PAGE;
-    return previews.slice(startIndex, startIndex + ITEMS_PER_PAGE);
-  }, [previews, currentPage]);
-
-  useEffect(() => {
-    const pageHasBeenRendered = renderedPages.has(currentPage);
-    const previewsOnThisPageNeedFetching = visiblePreviews.some(
-      (p) => !p.htmlFragment && !p.fragmentError
-    );
+  };
 
-    if (previewsOnThisPageNeedFetching && !pageHasBeenRendered) {
-      const newRequests = visiblePreviews
-        .filter((p) => !p.htmlFragment && !p.fragmentError)
-        .map((p) => ({
-          id: `template-${p.index}`,
-          ctx: p.ctx,
-        }));
-      setFragmentsToGenerate(newRequests);
-    } else {
-      setFragmentsToGenerate([]);
+  const handleBack = () => {
+    setError(null);
+    if (step === 'copyInput') {
+      setStep('initial');
+    } else if (step === 'designLibrary' || step === 'aiDesign' || 'error') {
+      setStep('copyInput');
     }
-  }, [currentPage, visiblePreviews, renderedPages]);
+  };
 
-  const handlePageChange = (newPage: number) => {
-    if (newPage >= 0 && newPage < totalPages) {
-      setCurrentPage(newPage);
+  const handleCopyContinue = () => {
+    if (initialChoice === 'library') {
+      setStep('designLibrary');
+    } else if (initialChoice === 'ai') {
+      setStep('aiDesign');
     }
   };
 
-  const handleFragmentsComplete = (results: PaneFragmentResult[]) => {
-    setPreviews((prevPreviews) => {
-      const updated = [...prevPreviews];
-      results.forEach((result) => {
-        const index = parseInt(result.id.replace('template-', ''));
-        const previewIndex = updated.findIndex((p) => p.index === index);
-        if (previewIndex !== -1) {
-          updated[previewIndex] = {
-            ...updated[previewIndex],
-            htmlFragment: result.htmlString,
-            fragmentError: result.error,
-          };
-        }
-      });
-      return updated;
-    });
-    setRenderedPages((prev) => new Set(prev).add(currentPage));
+  const handleBlankSlate = () => {
+    const blankTemplate: TemplatePane = {
+      id: '', // ctx will assign
+      nodeType: 'Pane',
+      parentId: '', // ctx will assign
+      title: 'New Pane',
+      slug: '',
+      isDecorative: false,
+      markdown: {
+        id: '', // ctx will assign
+        nodeType: 'Markdown',
+        parentId: '', // ctx will assign
+        type: 'markdown',
+        markdownId: '', // ctx will assign
+        defaultClasses: {},
+        parentClasses: [],
+        nodes: [],
+      },
+    };
+    handleApplyTemplate(blankTemplate);
   };
 
-  const handleSnapshotComplete = (id: string, snapshot: SnapshotData) => {
-    const index = parseInt(id.replace('template-', ''));
-    setPreviews((prevPreviews) => {
-      const updated = [...prevPreviews];
-      const previewIndex = updated.findIndex((p) => p.index === index);
-      if (previewIndex !== -1) {
-        updated[previewIndex] = {
-          ...updated[previewIndex],
-          snapshot,
-        };
+  const handleDesignLibrarySelect = async (entry: DesignLibraryEntry) => {
+    // This flow is for "Design Library + Provide Copy"
+    if (copyMode === 'raw') {
+      const liveTemplate = convertStorageToLiveTemplate(
+        mergeCopyIntoTemplate(entry.template, []) // Start with blank copy
+      );
+      if (liveTemplate.markdown) {
+        liveTemplate.markdown.markdownBody = copyValue;
       }
-      return updated;
-    });
-  };
-
-  const handleApplyTemplate = async (template: any) => {
-    if (isInserting || !ctx) return;
-    setIsInserting(true);
-
-    try {
-      const insertTemplate =
-        mode === 'custom'
-          ? {
-              ...cloneDeep(template),
-              markdown: template.markdown && {
-                ...template.markdown,
-                markdownBody: customMarkdown,
-              },
-            }
-          : cloneDeep(template);
-
-      insertTemplate.title = '';
-      insertTemplate.slug = '';
+      handleApplyTemplate(liveTemplate);
+      return;
+    }
 
-      const ownerId =
-        isStoryFragment || isContextPane
-          ? nodeId
-          : ctx.getClosestNodeTypeFromId(nodeId, 'StoryFragment');
+    // This flow is for "Design Library + Write a Prompt" (Hybrid AI)
+    if (copyMode === 'prompt') {
+      setError(null);
+      setStep('loading');
+      try {
+        // 1. Get the full, rich template from the library
+        const liveTemplate = convertStorageToLiveTemplate(entry.template);
+        if (!liveTemplate.markdown) {
+          throw new Error(
+            'The selected design library item is not compatible with this workflow as it has no markdown section.'
+          );
+        }
 
-      if (isContextPane) {
-        insertTemplate.isContextPane = true;
-        await ctx.applyAtomicUpdate(async (tmpCtx) => {
-          tmpCtx.addContextTemplatePane(ownerId, insertTemplate);
-        });
-      } else {
-        await ctx.applyAtomicUpdate(async (tmpCtx) => {
-          tmpCtx.addTemplatePane(
-            ownerId,
-            insertTemplate,
-            nodeId,
-            first ? 'before' : 'after'
+        // 2. Create the simplified shell for the AI
+        const shellJson = convertTemplateToAIShell(liveTemplate);
+        if (!shellJson || shellJson === '{}') {
+          throw new Error(
+            'Could not generate a valid AI shell from this design.'
           );
-        });
-        ctx.notifyNode(`root`);
+        }
+
+        // 3. Get the AI to write copy based on the shell and prompt
+        const copyPromptDetails = prompts.aiPaneCopyPrompt;
+        const layout = 'Text Only';
+        const formattedCopyPrompt = copyPromptDetails.user_template
+          .replace('{{COPY_INPUT}}', promptValue)
+          .replace(
+            '{{DESIGN_INPUT}}',
+            "N/A - Use the provided Shell JSON's design."
+          )
+          .replace('{{LAYOUT_TYPE}}', layout)
+          .replace('{{SHELL_JSON}}', shellJson);
+
+        const copyResult = await callAskLemurAPI(
+          formattedCopyPrompt,
+          copyPromptDetails.system || '',
+          false,
+          isSandboxMode
+        );
+
+        // 4. Parse ONLY the AI-generated HTML into content nodes
+        const newNodes = parseAiCopyHtml(copyResult, liveTemplate.markdown.id);
+
+        // 5. Create the final pane by cloning the original rich template
+        const finalPane = cloneDeep(liveTemplate);
+
+        // 6. Inject the new AI content, preserving the original rich design
+        finalPane.markdown!.nodes = newNodes;
+
+        // 7. Apply the complete, correctly merged pane
+        handleApplyTemplate(finalPane);
+      } catch (err: any) {
+        setError(err.message || 'Failed to generate AI copy for this design.');
+        setStep('error');
       }
-      setParentMode(PaneAddMode.DEFAULT, false);
-    } catch (error) {
-      console.error('Error inserting template:', error);
-    } finally {
-      setIsInserting(false);
     }
   };
 
-  const handleApplyGeneratedPane = async (pane: TemplatePane) => {
-    if (isInserting || !ctx) return;
-    setIsInserting(true);
+  const handleAiDesignGenerate = useCallback(async () => {
+    setError(null);
+    setStep('loading');
+
+    let designInput = `Generate a design using a **${aiDesignConfig.harmony.toLowerCase()}** color scheme with a **${aiDesignConfig.theme.toLowerCase()}** theme.`;
+    if (aiDesignConfig.baseColor)
+      designInput += ` Base the colors around **${aiDesignConfig.baseColor}**.`;
+    if (aiDesignConfig.accentColor)
+      designInput += ` Use **${aiDesignConfig.accentColor}** as an accent color.`;
+    if (aiDesignConfig.additionalNotes)
+      designInput += ` Refine with these notes: "${aiDesignConfig.additionalNotes}"`;
+
+    try {
+      const shellPromptDetails = prompts.aiPaneShellPrompt;
+      const copyPromptDetails = prompts.aiPaneCopyPrompt;
+      const layout = 'Text Only'; // Hardcoded for this simplified AI path
+
+      const formattedShellPrompt = shellPromptDetails.user_template
+        .replace('{{DESIGN_INPUT}}', designInput)
+        .replace('{{LAYOUT_TYPE}}', layout);
+
+      const shellResult = await callAskLemurAPI(
+        formattedShellPrompt,
+        shellPromptDetails.system || '',
+        true,
+        isSandboxMode
+      );
+
+      const copyInputContent = copyMode === 'prompt' ? promptValue : copyValue;
+      const formattedCopyPrompt = copyPromptDetails.user_template
+        .replace('{{COPY_INPUT}}', copyInputContent)
+        .replace('{{DESIGN_INPUT}}', designInput)
+        .replace('{{LAYOUT_TYPE}}', layout)
+        .replace('{{SHELL_JSON}}', shellResult);
+
+      const copyResult = await callAskLemurAPI(
+        formattedCopyPrompt,
+        copyPromptDetails.system || '',
+        false,
+        isSandboxMode
+      );
+
+      const finalPane = parseAiPane(shellResult, copyResult, layout);
+      handleApplyTemplate(finalPane);
+    } catch (err: any) {
+      setError(err.message || 'Failed to generate AI pane.');
+      setStep('error');
+    }
+  }, [aiDesignConfig, copyMode, promptValue, copyValue, isSandboxMode]);
+
+  const handleApplyTemplate = async (template: TemplatePane) => {
+    if (!ctx) return;
     try {
+      const insertTemplate = cloneDeep(template);
+      insertTemplate.title = insertTemplate.title || 'New Pane';
+      insertTemplate.slug = insertTemplate.slug || '';
+
       const ownerId =
         isStoryFragment || isContextPane
           ? nodeId
           : ctx.getClosestNodeTypeFromId(nodeId, 'StoryFragment');
 
       if (isContextPane) {
-        pane.isContextPane = true;
+        insertTemplate.isContextPane = true;
         await ctx.applyAtomicUpdate(async (tmpCtx) => {
-          tmpCtx.addContextTemplatePane(ownerId, pane);
+          tmpCtx.addContextTemplatePane(ownerId, insertTemplate);
         });
       } else {
         await ctx.applyAtomicUpdate(async (tmpCtx) => {
           tmpCtx.addTemplatePane(
             ownerId,
-            pane,
+            insertTemplate,
             nodeId,
             first ? 'before' : 'after'
           );
@@ -283,363 +345,203 @@ const AddPaneNewPanel = ({
         ctx.notifyNode(`root`);
       }
       setParentMode(PaneAddMode.DEFAULT, false);
-    } catch (error) {
-      console.error('Error applying generated pane:', error);
-    } finally {
-      setIsInserting(false);
+    } catch (err) {
+      console.error('Error inserting template:', err);
+      setError(
+        err instanceof Error
+          ? err.message
+          : 'An unexpected error occurred while adding the pane.'
+      );
+      setStep('error');
     }
   };
 
-  const handleThemeChange = (details: { value: string[] }) => {
-    const newTheme = details.value[0] as Theme;
-    if (newTheme) {
-      setSelectedTheme(newTheme);
-    }
-  };
-
-  const handleCategoryChange = (details: { value: string[] }) => {
-    const id = details.value[0];
-    if (id) {
-      const category = templateCategories.find((cat) => cat.id === id);
-      if (category) setSelectedCategory(category);
-    }
-  };
+  // --- Render Logic ---
 
-  const customStyles = `
-    .category-item[data-highlighted] { background-color: #0891b2; color: white; }
-    .category-item[data-highlighted] .category-indicator { color: white; }
-    .category-item[data-state="checked"] .category-indicator { display: flex; }
-    .category-item .category-indicator { display: none; }
-    .category-item[data-state="checked"] { font-weight: bold; }
-    .theme-item[data-highlighted] { background-color: #0891b2; color: white; }
-    .theme-item[data-highlighted] .theme-indicator { color: white; }
-    .theme-item[data-state="checked"] .theme-indicator { display: flex; }
-    .theme-item .theme-indicator { display: none; }
-    .theme-item[data-state="checked"] { font-weight: bold; }
-  `;
+  const renderInitialStep = () => (
+    <div className="p-4">
+      <h3 className="font-action mb-4 text-center text-xl font-bold text-gray-800">
+        How would you like to design your new pane?
+      </h3>
+      <div className="grid grid-cols-1 gap-4 md:grid-cols-3">
+        <button
+          onClick={() => handleInitialChoice('library')}
+          className="group flex flex-col items-center space-y-3 rounded-lg border bg-white p-6 text-center shadow-sm transition-all hover:border-cyan-600 hover:shadow-lg"
+        >
+          <SwatchIcon className="h-10 w-10 text-gray-500 transition-colors group-hover:text-cyan-600" />
+          <h4 className="font-semibold text-gray-800">Use Design Library</h4>
+          <p className="text-sm text-gray-600">
+            Start with a pre-made design and add your own content.
+          </p>
+        </button>
+        {hasAssemblyAI && (
+          <button
+            onClick={() => handleInitialChoice('ai')}
+            className="group flex flex-col items-center space-y-3 rounded-lg border bg-white p-6 text-center shadow-sm transition-all hover:border-cyan-600 hover:shadow-lg"
+          >
+            <SparklesIcon className="h-10 w-10 text-gray-500 transition-colors group-hover:text-cyan-600" />
+            <h4 className="font-semibold text-gray-800">Design with AI</h4>
+            <p className="text-sm text-gray-600">
+              Let AI generate a complete design and copy from your prompt.
+            </p>
+          </button>
+        )}
+        <button
+          onClick={() => handleInitialChoice('blank')}
+          className="group flex flex-col items-center space-y-3 rounded-lg border bg-white p-6 text-center shadow-sm transition-all hover:border-cyan-600 hover:shadow-lg"
+        >
+          <DocumentPlusIcon className="h-10 w-10 text-gray-500 transition-colors group-hover:text-cyan-600" />
+          <h4 className="font-semibold text-gray-800">Blank Slate</h4>
+          <p className="text-sm text-gray-600">
+            Add a simple, empty pane to build from scratch.
+          </p>
+        </button>
+      </div>
+    </div>
+  );
 
-  return (
-    <div className="bg-white p-3.5 shadow-inner">
-      <style>{customStyles}</style>
-      <div className="group flex w-full gap-1 rounded-md bg-white p-1.5">
+  const renderContentStep = () => (
+    <div className="space-y-4 p-4">
+      <CopyInputStep
+        copyMode={copyMode}
+        onCopyModeChange={setCopyMode}
+        promptValue={promptValue}
+        onPromptValueChange={setPromptValue}
+        copyValue={copyValue}
+        onCopyValueChange={setCopyValue}
+        defaultPrompt={
+          first
+            ? prompts.aiPaneCopyPrompt.heroDefault
+            : prompts.aiPaneCopyPrompt.contentDefault
+        }
+      />
+      <div className="flex justify-between">
         <button
-          onClick={() => setParentMode(PaneAddMode.DEFAULT, first)}
-          className="w-fit rounded bg-gray-100 px-3 py-1 text-sm text-gray-700 transition-colors hover:bg-gray-200 focus:bg-gray-200"
-          type="button"
+          onClick={handleBack}
+          className="rounded-md border border-gray-300 bg-white px-4 py-2 text-sm font-bold text-gray-700 shadow-sm hover:bg-gray-50"
         >
-          ← Go Back
+          ← Back
         </button>
+        <button
+          onClick={handleCopyContinue}
+          disabled={
+            copyMode === 'prompt' ? !promptValue.trim() : !copyValue.trim()
+          }
+          className="rounded-md bg-cyan-600 px-4 py-2 text-sm font-bold text-white shadow-sm hover:bg-cyan-700 disabled:cursor-not-allowed disabled:bg-gray-400"
+        >
+          Continue →
+        </button>
+      </div>
+    </div>
+  );
 
-        <div className="ml-4 flex flex-wrap items-center gap-x-6 gap-y-2 py-2">
-          <div className="font-action flex-none rounded px-2 py-2.5 text-sm font-bold text-cyan-700 shadow-sm">
-            + Design New Pane
-          </div>
-          <div className="flex items-center space-x-2 rounded-lg bg-gray-100 p-1">
-            <button
-              onClick={() => setMode('template')}
-              className={`rounded-md px-3 py-1 text-sm font-medium transition-colors ${
-                mode === 'template'
-                  ? 'bg-white text-cyan-700 shadow'
-                  : 'text-gray-600 hover:text-gray-800'
-              }`}
-              type="button"
-            >
-              Use Template
-            </button>
-            <button
-              onClick={() => setMode('custom')}
-              className={`rounded-md px-3 py-1 text-sm font-medium transition-colors ${
-                mode === 'custom'
-                  ? 'bg-white text-cyan-700 shadow'
-                  : 'text-gray-600 hover:text-gray-800'
-              }`}
-              type="button"
-            >
-              Paste Markdown
-            </button>
-            {hasAssemblyAI && (
-              <button
-                onClick={() => setMode('ai')}
-                className={`rounded-md px-3 py-1 text-sm font-medium transition-colors ${
-                  mode === 'ai'
-                    ? 'rounded-md border border-transparent bg-cyan-600 px-3 py-1.5 text-sm font-bold text-white shadow-sm hover:bg-cyan-700 focus:outline-none focus:ring-2 focus:ring-cyan-500 focus:ring-offset-2'
-                    : 'text-gray-600 hover:text-gray-800'
-                }`}
-                type="button"
-              >
-                ✨ Generate with AI
-              </button>
-            )}
-          </div>
+  const renderDesignLibraryStep = () => (
+    <div className="space-y-4 p-4">
+      <div className="flex justify-start">
+        <button
+          onClick={handleBack}
+          className="rounded-md border border-gray-300 bg-white px-4 py-2 text-sm font-bold text-gray-700 shadow-sm hover:bg-gray-50"
+        >
+          ← Back to Content
+        </button>
+      </div>
+      <DesignLibraryStep
+        config={config!}
+        onSelect={handleDesignLibrarySelect}
+      />
+    </div>
+  );
 
-          {mode === 'custom' && (
-            <div className="mt-4 w-full">
-              <AddPaneNewCustomCopy
-                value={customMarkdown}
-                onChange={setCustomMarkdown}
-              />
-            </div>
-          )}
-          {mode === 'ai' && (
-            <div className="mt-4 w-full">
-              <AiPaneGenerator
-                ownerId={nodeId}
-                onComplete={handleApplyGeneratedPane}
-                onCancel={() => setMode('template')}
-                config={config!}
-              />
-            </div>
-          )}
-        </div>
+  const renderAiDesignStep = () => (
+    <div className="space-y-4 p-4">
+      <AiDesignStep
+        config={config!}
+        designConfig={aiDesignConfig}
+        onDesignConfigChange={setAiDesignConfig}
+      />
+      <div className="flex justify-between">
+        <button
+          onClick={handleBack}
+          className="rounded-md border border-gray-300 bg-white px-4 py-2 text-sm font-bold text-gray-700 shadow-sm hover:bg-gray-50"
+        >
+          ← Back
+        </button>
+        <button
+          onClick={handleAiDesignGenerate}
+          className="rounded-md bg-cyan-600 px-4 py-2 text-sm font-bold text-white shadow-sm hover:bg-cyan-700"
+        >
+          ✨ Generate with AI
+        </button>
       </div>
+    </div>
+  );
 
-      {mode !== 'ai' && (
-        <>
-          <h3 className="font-action px-3.5 pb-1.5 pt-4 text-xl font-bold text-black">
-            1. Template design settings
-          </h3>
-
-          <div className="grid grid-cols-1 gap-4 p-2 md:grid-cols-3">
-            <div className="w-full">
-              <Select.Root
-                collection={themesCollection}
-                value={[selectedTheme]}
-                onValueChange={handleThemeChange}
-              >
-                <Select.Label className="block text-sm font-bold text-gray-700">
-                  Theme
-                </Select.Label>
-                <Select.Control className="relative mt-1">
-                  <Select.Trigger className="relative w-full cursor-default rounded-lg bg-white py-2 pl-3 pr-10 text-left shadow-md focus:outline-none focus-visible:border-cyan-600 focus-visible:ring-2 focus-visible:ring-white/75 focus-visible:ring-offset-2 focus-visible:ring-offset-cyan-600">
-                    <Select.ValueText className="block truncate capitalize">
-                      {selectedTheme.replace(/-/g, ' ')}
-                    </Select.ValueText>
-                    <Select.Indicator className="absolute inset-y-0 right-0 flex items-center pr-2">
-                      <ChevronUpDownIcon
-                        className="h-5 w-5 text-gray-400"
-                        aria-hidden="true"
-                      />
-                    </Select.Indicator>
-                  </Select.Trigger>
-                </Select.Control>
-                <Portal>
-                  <Select.Positioner>
-                    <Select.Content className="z-50 max-h-60 w-full overflow-auto rounded-md bg-white py-1 text-base shadow-lg ring-1 ring-black/5 focus:outline-none sm:text-sm">
-                      {themesCollection.items.map((theme) => (
-                        <Select.Item
-                          key={theme}
-                          item={theme}
-                          className="theme-item relative cursor-default select-none py-2 pl-10 pr-4 text-gray-900"
-                        >
-                          <Select.ItemText className="block truncate capitalize">
-                            {theme.replace(/-/g, ' ')}
-                          </Select.ItemText>
-                          <Select.ItemIndicator className="theme-indicator absolute inset-y-0 left-0 flex items-center pl-3 text-cyan-600">
-                            <CheckIcon className="h-5 w-5" aria-hidden="true" />
-                          </Select.ItemIndicator>
-                        </Select.Item>
-                      ))}
-                    </Select.Content>
-                  </Select.Positioner>
-                </Portal>
-              </Select.Root>
-            </div>
+  const renderLoading = () => (
+    <div className="flex min-h-[300px] flex-col items-center justify-center space-y-4 p-6">
+      <div className="h-8 w-8 animate-spin rounded-full border-b-2 border-cyan-600"></div>
+      <p className="text-sm text-gray-600">Generating AI Pane...</p>
+      <p className="text-xs text-gray-500">This may take a moment.</p>
+    </div>
+  );
 
-            {!isContextPane && (
-              <div className="w-full">
-                <Select.Root
-                  collection={categoryCollection}
-                  value={[selectedCategory.id]}
-                  onValueChange={handleCategoryChange}
-                >
-                  <Select.Label className="block text-sm font-bold text-gray-700">
-                    Category
-                  </Select.Label>
-                  <Select.Control className="relative mt-1">
-                    <Select.Trigger className="relative w-full cursor-default rounded-lg bg-white py-2 pl-3 pr-10 text-left shadow-md focus:outline-none focus-visible:ring-2 focus-visible:ring-white/75 focus-visible:ring-offset-2 focus-visible:ring-offset-teal-300 sm:text-sm">
-                      <Select.ValueText className="block truncate">
-                        {selectedCategory.title}
-                      </Select.ValueText>
-                      <Select.Indicator className="absolute inset-y-0 right-0 flex items-center pr-2">
-                        <ChevronUpDownIcon
-                          className="h-5 w-5 text-gray-400"
-                          aria-hidden="true"
-                        />
-                      </Select.Indicator>
-                    </Select.Trigger>
-                  </Select.Control>
-                  <Portal>
-                    <Select.Positioner>
-                      <Select.Content className="z-50 max-h-60 w-full overflow-auto rounded-md bg-white py-1 text-base shadow-lg ring-1 ring-black/5 focus:outline-none sm:text-sm">
-                        {categoryCollection.items.map((category) => (
-                          <Select.Item
-                            key={category.id}
-                            item={category}
-                            className="category-item relative cursor-default select-none py-2 pl-10 pr-4 text-gray-900"
-                          >
-                            <Select.ItemText className="block truncate">
-                              {category.title}
-                            </Select.ItemText>
-                            <Select.ItemIndicator className="category-indicator absolute inset-y-0 left-0 flex items-center pl-3 text-cyan-600">
-                              <CheckIcon
-                                className="h-5 w-5"
-                                aria-hidden="true"
-                              />
-                            </Select.ItemIndicator>
-                          </Select.Item>
-                        ))}
-                      </Select.Content>
-                    </Select.Positioner>
-                  </Portal>
-                </Select.Root>
-              </div>
-            )}
-
-            <div className="flex items-center gap-2">
-              <Switch.Root
-                checked={useOddVariant}
-                onCheckedChange={(details) => setUseOddVariant(details.checked)}
-                className="inline-flex items-center"
-              >
-                <Switch.Control
-                  className={`${
-                    useOddVariant ? 'bg-cyan-600' : 'bg-gray-200'
-                  } relative my-2 inline-flex h-6 w-11 items-center rounded-full transition-colors focus:outline-none focus:ring-2 focus:ring-cyan-500 focus:ring-offset-2`}
-                >
-                  <Switch.Thumb
-                    className={`${
-                      useOddVariant ? 'translate-x-6' : 'translate-x-1'
-                    } inline-block h-4 w-4 rounded-full bg-white shadow-lg transition-transform duration-200`}
-                  />
-                </Switch.Control>
-                <Switch.HiddenInput />
-                <div className="flex h-6 items-center">
-                  <Switch.Label className="px-4 text-sm text-gray-700">
-                    Toggle subtle variation (great for stacking content)
-                  </Switch.Label>
-                </div>
-              </Switch.Root>
-            </div>
-          </div>
+  const renderError = () => (
+    <div className="space-y-4 rounded-lg bg-red-50 p-6 text-center">
+      <h4 className="text-lg font-bold text-red-800">Generation Failed</h4>
+      <p className="text-sm text-red-700">{error}</p>
+      <button
+        onClick={handleBack}
+        className="mt-2 rounded-md border border-red-300 bg-white px-4 py-2 text-sm font-bold text-red-800 shadow-sm hover:bg-red-50"
+      >
+        ← Try Again
+      </button>
+    </div>
+  );
 
-          <h3 className="font-action px-3.5 pb-1.5 pt-4 text-xl font-bold text-black">
-            2. Choose design
-          </h3>
-
-          {fragmentsToGenerate.length > 0 && (
-            <PanesPreviewGenerator
-              requests={fragmentsToGenerate}
-              onComplete={handleFragmentsComplete}
-              onError={(error) =>
-                console.error('Fragment generation error:', error)
-              }
-            />
-          )}
-
-          <div className="grid grid-cols-2 gap-4 p-2 xl:grid-cols-3">
-            {visiblePreviews.map((preview) => (
-              <div key={preview.index} className="flex flex-col items-center">
-                <div
-                  onClick={
-                    isInserting
-                      ? undefined
-                      : () => handleApplyTemplate(preview.template)
-                  }
-                  className={`bg-mywhite group relative w-full rounded-sm shadow-inner ${
-                    isInserting
-                      ? 'cursor-not-allowed opacity-50'
-                      : 'cursor-pointer'
-                  } transition-all duration-200 ${
-                    preview.snapshot
-                      ? 'hover:outline-solid hover:outline hover:outline-4'
-                      : ''
-                  }`}
-                  style={{
-                    ...(!preview.snapshot ? { minHeight: '150px' } : {}),
-                  }}
-                  role="button"
-                  tabIndex={0}
-                  aria-label={preview.template.title}
-                >
-                  {preview.htmlFragment &&
-                    !preview.snapshot &&
-                    !preview.fragmentError && (
-                      <PaneSnapshotGenerator
-                        id={`template-${preview.index}`}
-                        htmlString={preview.htmlFragment}
-                        onComplete={handleSnapshotComplete}
-                        onError={(id, error) =>
-                          console.error(`Snapshot error for ${id}:`, error)
-                        }
-                        outputWidth={800}
-                      />
-                    )}
-
-                  {!preview.htmlFragment && !preview.fragmentError && (
-                    <div className="flex h-48 items-center justify-center">
-                      <div className="text-gray-500">Loading preview...</div>
-                    </div>
-                  )}
-
-                  {preview.fragmentError && (
-                    <div className="flex h-48 items-center justify-center">
-                      <div className="text-red-500">Preview error</div>
-                    </div>
-                  )}
-
-                  {preview.snapshot && (
-                    <div className="p-0.5">
-                      <img
-                        src={preview.snapshot.imageData}
-                        alt={`Template: ${preview.template.title}`}
-                        className="w-full"
-                      />
-                    </div>
-                  )}
-                </div>
-                <p className="bg-mydarkgrey mt-2 w-full break-words p-2 text-center text-sm text-white">
-                  {preview.template.title}
-                </p>
-              </div>
-            ))}
-          </div>
+  const renderStep = () => {
+    switch (step) {
+      case 'initial':
+        return renderInitialStep();
+      case 'copyInput':
+        return renderContentStep();
+      case 'designLibrary':
+        return renderDesignLibraryStep();
+      case 'aiDesign':
+        return renderAiDesignStep();
+      case 'loading':
+        return renderLoading();
+      case 'error':
+        return renderError();
+      default:
+        return renderInitialStep();
+    }
+  };
 
-          <div className="mb-2 mt-4 flex items-center justify-center gap-2">
+  return (
+    <div className="bg-white p-2 shadow-inner">
+      <div className="group mb-2 flex w-full items-center gap-1 rounded-md bg-white p-1.5">
+        {first ? (
+          <div className="w-full text-center">
+            <h2 className="font-action py-1.5 text-lg font-bold text-gray-800">
+              Welcome to Tract Stack
+            </h2>
+          </div>
+        ) : (
+          <>
             <button
-              onClick={() => handlePageChange(currentPage - 1)}
-              disabled={currentPage === 0}
-              className="rounded bg-gray-100 px-3 py-1 text-sm text-gray-700 transition-colors hover:bg-gray-200 focus:bg-gray-200 disabled:cursor-not-allowed disabled:opacity-50"
-              type="button"
+              onClick={() => setParentMode(PaneAddMode.DEFAULT, first)}
+              className="w-fit rounded bg-gray-100 px-3 py-1 text-sm text-gray-700 transition-colors hover:bg-gray-200"
             >
-              Previous
+              ← Go Back
             </button>
-            <div className="flex gap-1">
-              {[...Array(totalPages)].map((_, index) => (
-                <button
-                  key={index}
-                  onClick={() => handlePageChange(index)}
-                  className={`rounded px-3 py-1 text-sm transition-colors ${
-                    currentPage === index
-                      ? 'bg-cyan-700 text-white'
-                      : 'bg-gray-100 text-gray-700 hover:bg-gray-200'
-                  }`}
-                  type="button"
-                >
-                  {index + 1}
-                </button>
-              ))}
+            <div className="font-action ml-4 flex-none rounded px-2 py-2.5 text-sm font-bold text-cyan-700 shadow-sm">
+              + Design New Pane
             </div>
-            <button
-              onClick={() => handlePageChange(currentPage + 1)}
-              disabled={currentPage === totalPages - 1}
-              className="rounded bg-gray-100 px-3 py-1 text-sm text-gray-700 transition-colors hover:bg-gray-200 focus:bg-gray-200 disabled:cursor-not-allowed disabled:opacity-50"
-              type="button"
-            >
-              Next
-            </button>
-          </div>
-        </>
-      )}
+          </>
+        )}
+      </div>
+      <div className="min-h-[400px] rounded-md border bg-gray-50">
+        {renderStep()}
+      </div>
     </div>
   );
 };