"use client"; import { useEffect, useRef, useState, useCallback } from "react"; import { NodeType } from "@/types"; // Actions are special menu items that trigger behavior instead of creating a node export type MenuAction = "splitGridImmediate"; interface MenuOption { type: NodeType | MenuAction; label: string; icon: React.ReactNode; isAction?: boolean; // true if this is an action, not a node type } // Define which nodes can accept which handle types as inputs const IMAGE_TARGET_OPTIONS: MenuOption[] = [ { type: "annotation", label: "Annotate", icon: ( ), }, { type: "nanoBanana", label: "Generate Image", icon: ( ), }, { type: "generateVideo", label: "Generate Video", icon: ( ), }, { type: "splitGrid", label: "Split Grid Node", icon: ( ), }, { type: "splitGridImmediate", label: "Split Grid Now", isAction: true, icon: ( ), }, { type: "output", label: "Output", icon: ( ), }, { type: "outputGallery", label: "Output Gallery", icon: ( ), }, { type: "imageCompare", label: "Image Compare", icon: ( ), }, ]; const TEXT_TARGET_OPTIONS: MenuOption[] = [ { type: "prompt", label: "Prompt", icon: ( ), }, { type: "promptConstructor", label: "Prompt Constructor", icon: ( ), }, { type: "nanoBanana", label: "Generate Image", icon: ( ), }, { type: "generateVideo", label: "Generate Video", icon: ( ), }, { type: "llmGenerate", label: "LLM Generate", icon: ( ), }, ]; // Define which nodes can provide sources for handle types (when dragging to a target handle) const IMAGE_SOURCE_OPTIONS: MenuOption[] = [ { type: "imageInput", label: "Image Input", icon: ( ), }, { type: "annotation", label: "Annotate", icon: ( ), }, { type: "nanoBanana", label: "Generate Image", icon: ( ), }, ]; const TEXT_SOURCE_OPTIONS: MenuOption[] = [ { type: "prompt", label: "Prompt", icon: ( ), }, { type: "promptConstructor", label: "Prompt Constructor", icon: ( ), }, { type: "llmGenerate", label: "LLM Generate", icon: ( ), }, ]; // Video can only connect to generateVideo (video-to-video) or output nodes const VIDEO_TARGET_OPTIONS: MenuOption[] = [ { type: "generateVideo", label: "Generate Video", icon: ( ), }, { type: "output", label: "Output", icon: ( ), }, ]; // Only generateVideo nodes produce video output const VIDEO_SOURCE_OPTIONS: MenuOption[] = [ { type: "generateVideo", label: "Generate Video", icon: ( ), }, ]; // Audio target options (nodes that accept audio input) const AUDIO_TARGET_OPTIONS: MenuOption[] = [ // VideoStitch will be added in Plan 03 ]; // Audio source options (nodes that produce audio output) const AUDIO_SOURCE_OPTIONS: MenuOption[] = [ { type: "audioInput", label: "Audio Input", icon: ( ), }, ]; interface ConnectionDropMenuProps { position: { x: number; y: number }; handleType: "image" | "text" | "video" | "audio" | null; connectionType: "source" | "target"; // source = dragging from output, target = dragging from input onSelect: (selection: { type: NodeType | MenuAction; isAction: boolean }) => void; onClose: () => void; } export function ConnectionDropMenu({ position, handleType, connectionType, onSelect, onClose, }: ConnectionDropMenuProps) { const menuRef = useRef(null); const [selectedIndex, setSelectedIndex] = useState(0); // Get the appropriate node options based on handle type and connection direction const getOptions = useCallback((): MenuOption[] => { if (!handleType) return []; if (connectionType === "source") { // Dragging from a source handle (output), need nodes with target handles (inputs) if (handleType === "video") return VIDEO_TARGET_OPTIONS; if (handleType === "audio") return AUDIO_TARGET_OPTIONS; return handleType === "image" ? IMAGE_TARGET_OPTIONS : TEXT_TARGET_OPTIONS; } else { // Dragging from a target handle (input), need nodes with source handles (outputs) if (handleType === "video") return VIDEO_SOURCE_OPTIONS; if (handleType === "audio") return AUDIO_SOURCE_OPTIONS; return handleType === "image" ? IMAGE_SOURCE_OPTIONS : TEXT_SOURCE_OPTIONS; } }, [handleType, connectionType]); const options = getOptions(); // Handle keyboard navigation useEffect(() => { const handleKeyDown = (e: KeyboardEvent) => { switch (e.key) { case "ArrowDown": e.preventDefault(); setSelectedIndex((prev) => (prev + 1) % options.length); break; case "ArrowUp": e.preventDefault(); setSelectedIndex((prev) => (prev - 1 + options.length) % options.length); break; case "Enter": e.preventDefault(); if (options[selectedIndex]) { onSelect({ type: options[selectedIndex].type, isAction: options[selectedIndex].isAction || false, }); } break; case "Escape": e.preventDefault(); onClose(); break; } }; document.addEventListener("keydown", handleKeyDown); return () => document.removeEventListener("keydown", handleKeyDown); }, [options, selectedIndex, onSelect, onClose]); // Close when clicking outside useEffect(() => { const handleClickOutside = (e: MouseEvent) => { if (menuRef.current && !menuRef.current.contains(e.target as Node)) { onClose(); } }; document.addEventListener("mousedown", handleClickOutside); return () => document.removeEventListener("mousedown", handleClickOutside); }, [onClose]); // Focus the menu when it opens useEffect(() => { menuRef.current?.focus(); }, []); if (options.length === 0) return null; return (
Add {handleType} node
{options.map((option, index) => ( ))}
↑↓ navigate ↵ select
); }