import { useMemo, useCallback, useRef, useState, useEffect } from 'react'; import ForceGraph3D from 'react-force-graph-3d'; import { useConversation } from '../context/ConversationContext'; import { toGraphData } from '../utils/tree'; import BranchDialog from './BranchDialog'; import * as THREE from 'three'; export default function TreeView() { const { state, switchToBranch, setNodeAsMainBranch, } = useConversation(); const fgRef = useRef(null); const containerRef = useRef(null); const treePaneRef = useRef(null); const prevNodeCountRef = useRef(0); const [dimensions, setDimensions] = useState({ width: 420, height: 600 }); const [contextMenu, setContextMenu] = useState(null); const [isFullscreen, setIsFullscreen] = useState(false); const [branchNode, setBranchNode] = useState(null); const [showLegend, setShowLegend] = useState(true); const conv = state.activeConversationId ? state.conversations[state.activeConversationId] : null; // --- Stable graph structure fingerprint --- const structureFingerprint = useMemo(() => { if (!conv) return ''; const nodeIds = Object.keys(conv.nodes).sort().join(','); const links = Object.values(conv.nodes) .filter((n) => n.parentId) .map((n) => `${n.parentId}->${n.id}`) .sort() .join(','); const mainPath = conv.mainBranchPath.join(','); const activePath = conv.activeBranchPath.join(','); const statuses = Object.values(conv.nodes) .map((n) => `${n.id}:${n.status}`) .join(','); return `${nodeIds}|${links}|${mainPath}|${activePath}|${statuses}`; }, [conv]); const graphData = useMemo(() => { if (!conv || Object.keys(conv.nodes).length === 0) { return { nodes: [], links: [] }; } return toGraphData(conv.nodes, conv.mainBranchPath, conv.activeBranchPath); // eslint-disable-next-line react-hooks/exhaustive-deps }, [structureFingerprint]); // Auto-fit camera when new nodes are added useEffect(() => { const currentCount = graphData.nodes.length; if (currentCount > prevNodeCountRef.current && currentCount > 0 && fgRef.current) { setTimeout(() => { fgRef.current?.zoomToFit(600, 60); }, 300); } prevNodeCountRef.current = currentCount; }, [graphData.nodes.length]); // Resize observer useEffect(() => { if (!containerRef.current) return; const observer = new ResizeObserver((entries) => { const { width, height } = entries[0].contentRect; setDimensions({ width, height }); }); observer.observe(containerRef.current); return () => observer.disconnect(); }, []); // Close context menu on click outside useEffect(() => { if (!contextMenu) return; const handler = () => setContextMenu(null); window.addEventListener('click', handler); return () => window.removeEventListener('click', handler); }, [contextMenu]); // Handle Escape to exit fullscreen useEffect(() => { if (!isFullscreen) return; const handler = (e) => { if (e.key === 'Escape') setIsFullscreen(false); }; window.addEventListener('keydown', handler); return () => window.removeEventListener('keydown', handler); }, [isFullscreen]); const handleNodeClick = useCallback( (node) => { if (!conv) return; switchToBranch(state.activeConversationId, node.id); }, [conv, state.activeConversationId, switchToBranch] ); const handleNodeRightClick = useCallback( (node, event) => { event.preventDefault(); setContextMenu({ x: event.clientX, y: event.clientY, nodeId: node.id, userMessage: node.userMessage || '', }); }, [] ); const handleSetMainBranch = useCallback(() => { if (!contextMenu || !conv) return; setNodeAsMainBranch(state.activeConversationId, contextMenu.nodeId); setContextMenu(null); }, [contextMenu, conv, state.activeConversationId, setNodeAsMainBranch]); const handleBranchFromGraph = useCallback(() => { if (!contextMenu || !conv) return; const node = conv.nodes[contextMenu.nodeId]; if (!node) return; setBranchNode(node); setContextMenu(null); }, [contextMenu, conv]); const handleResetView = useCallback(() => { if (!fgRef.current) return; fgRef.current.zoomToFit(400, 40); }, []); const toggleFullscreen = useCallback(() => { setIsFullscreen((prev) => !prev); setTimeout(() => { if (fgRef.current) { fgRef.current.zoomToFit(400, 40); } }, 100); }, []); // --- Color helpers --- const getNodeColor = useCallback((node) => { if (node.status === 'generating') return '#f0c040'; if (node.isActive && node.isMain) return '#00ffcc'; if (node.isActive) return '#6c7aff'; if (node.isMain) return '#00d4aa'; return '#555566'; }, []); const linkColor = useCallback((link) => { if (link.isActive && link.isMain) return 'rgba(0, 255, 204, 0.5)'; if (link.isActive) return 'rgba(108, 122, 255, 0.4)'; if (link.isMain) return 'rgba(0, 212, 170, 0.35)'; return 'rgba(85, 85, 102, 0.2)'; }, []); const linkWidth = useCallback((link) => { if (link.isMain || link.isActive) return 2.5; return 1; }, []); // --- Rich hover tooltip --- const nodeLabel = useCallback((node) => { const userPreview = node.userMessage ? node.userMessage.slice(0, 120) + (node.userMessage.length > 120 ? '…' : '') : '(start)'; const aiPreview = node.assistantMessage ? node.assistantMessage.slice(0, 200) + (node.assistantMessage.length > 200 ? '…' : '') : node.status === 'generating' ? '⏳ Generating…' : '—'; const statusBadge = node.status === 'generating' ? 'GENERATING' : node.childCount > 1 ? `${node.childCount} branches` : ''; return `
Start chatting to see your conversation tree