src/components/TreeView.jsx

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'
        ? '<span style="background:#f0c040;color:#12121a;padding:2px 6px;border-radius:4px;font-size:9px;font-weight:700;margin-left:6px;">GENERATING</span>'
        : node.childCount > 1
        ? `<span style="background:rgba(108,122,255,0.2);color:#6c7aff;padding:2px 6px;border-radius:4px;font-size:9px;font-weight:600;margin-left:6px;">${node.childCount} branches</span>`
        : '';

    return `
      <div style="background:rgba(18,18,26,0.95);backdrop-filter:blur(16px);border:1px solid rgba(255,255,255,0.08);border-radius:12px;padding:14px 18px;max-width:360px;box-shadow:0 12px 40px rgba(0,0,0,0.5);font-family:Inter,sans-serif;">
        <div style="display:flex;align-items:center;margin-bottom:8px;">
          <span style="font-size:10px;text-transform:uppercase;letter-spacing:0.8px;font-weight:700;color:#6c7aff;">Turn ${node.depth + 1}</span>
          ${statusBadge}
        </div>
        <div style="font-size:10px;text-transform:uppercase;letter-spacing:0.8px;font-weight:600;color:rgba(108,122,255,0.6);margin-bottom:3px;">👤 User</div>
        <div style="font-size:12px;color:#e8e8f0;margin-bottom:12px;line-height:1.6;word-break:break-word;">${escapeHtml(userPreview)}</div>
        <div style="font-size:10px;text-transform:uppercase;letter-spacing:0.8px;font-weight:600;color:rgba(0,212,170,0.6);margin-bottom:3px;">🤖 Assistant</div>
        <div style="font-size:12px;color:#9898a8;line-height:1.6;word-break:break-word;">${escapeHtml(aiPreview)}</div>
      </div>
    `;
  }, []);

  // --- Custom 3D node with text label ---
  const nodeThreeObject = useCallback((node) => {
    const group = new THREE.Group();

    // Core sphere
    const isHighlighted = node.isMain || node.isActive;
    const size = isHighlighted ? 5 : 3;
    const color = getNodeColor(node);
    const geo = new THREE.SphereGeometry(size, 20, 20);
    const mat = new THREE.MeshPhongMaterial({
      color,
      emissive: color,
      emissiveIntensity: node.isActive ? 0.6 : 0.25,
      transparent: true,
      opacity: node.status === 'generating' ? 0.65 : 1,
      shininess: 60,
    });
    const mesh = new THREE.Mesh(geo, mat);
    group.add(mesh);

    // Outer glow ring for highlighted nodes
    if (isHighlighted) {
      const ringGeo = new THREE.RingGeometry(size * 1.4, size * 1.7, 32);
      const ringMat = new THREE.MeshBasicMaterial({
        color,
        transparent: true,
        opacity: 0.18,
        side: THREE.DoubleSide,
      });
      const ringMesh = new THREE.Mesh(ringGeo, ringMat);
      group.add(ringMesh);
    }

    // Branch indicator (diamond) for nodes with multiple children
    if (node.childCount > 1) {
      const diamondGeo = new THREE.OctahedronGeometry(2.5);
      const diamondMat = new THREE.MeshPhongMaterial({
        color: '#ff6c8a',
        emissive: '#ff6c8a',
        emissiveIntensity: 0.3,
        transparent: true,
        opacity: 0.7,
      });
      const diamond = new THREE.Mesh(diamondGeo, diamondMat);
      diamond.position.set(size + 4, 0, 0);
      diamond.rotation.set(0, 0, Math.PI / 4);
      group.add(diamond);
    }

    // Sprite label below the node (pure WebGL, no DOM elements)
    const label = createSpriteLabel(node, color);
    label.position.set(0, -(size + 8), 0);
    group.add(label);

    return group;
  }, [getNodeColor]);

  const nodeCount = graphData.nodes.length;

  return (
    <div
      className={`tree-pane ${isFullscreen ? 'tree-pane-fullscreen' : ''}`}
      ref={treePaneRef}
    >
      <div className="tree-header">
        <span>Graph View</span>
        <div className="tree-header-actions">
          <span style={{ fontSize: '11px', color: 'var(--text-tertiary)', marginRight: '8px' }}>
            {nodeCount} node{nodeCount !== 1 ? 's' : ''}
          </span>
          {nodeCount > 0 && (
            <>
              <button
                className="tree-action-btn"
                onClick={() => setShowLegend((v) => !v)}
                title="Toggle legend"
              >
                ℹ
              </button>
              <button
                className="tree-action-btn"
                onClick={handleResetView}
                title="Reset view (zoom to fit)"
              >
                ⟳
              </button>
            </>
          )}
          <button
            className="tree-action-btn"
            onClick={toggleFullscreen}
            title={isFullscreen ? 'Exit fullscreen (Esc)' : 'Fullscreen'}
          >
            {isFullscreen ? '⊗' : '⛶'}
          </button>
        </div>
      </div>

      <div className="tree-container" ref={containerRef}>
        {nodeCount === 0 ? (
          <div className="tree-empty">
            <p>Start chatting to see your conversation tree</p>
          </div>
        ) : (
          <ForceGraph3D
            ref={fgRef}
            graphData={graphData}
            width={dimensions.width}
            height={dimensions.height}
            backgroundColor="#0d0d14"
            dagMode="td"
            dagLevelDistance={80}
            nodeLabel={nodeLabel}
            nodeColor={getNodeColor}
            nodeVal={(node) => (node.isMain || node.isActive ? 3 : 1.5)}
            nodeThreeObject={nodeThreeObject}
            nodeThreeObjectExtend={false}
            linkColor={linkColor}
            linkWidth={linkWidth}
            linkCurvature={0.15}
            linkCurveRotation={0}
            linkDirectionalArrowLength={4}
            linkDirectionalArrowRelPos={0.85}
            linkDirectionalArrowColor={linkColor}
            linkDirectionalParticles={(link) =>
              link.isActive ? 3 : link.isMain ? 2 : 0
            }
            linkDirectionalParticleWidth={(link) =>
              link.isActive ? 2.5 : 1.5
            }
            linkDirectionalParticleColor={(link) =>
              link.isActive ? '#6c7aff' : '#00d4aa'
            }
            linkDirectionalParticleSpeed={0.004}
            linkOpacity={0.7}
            onNodeClick={handleNodeClick}
            onNodeRightClick={handleNodeRightClick}
            enableNodeDrag={false}
            cooldownTicks={100}
            warmupTicks={50}
            d3AlphaDecay={0.08}
            d3VelocityDecay={0.4}
            showNavInfo={false}
          />
        )}

        {/* Color Legend */}
        {showLegend && nodeCount > 0 && (
          <div className="graph-legend">
            <div className="graph-legend-title">Legend</div>
            <div className="graph-legend-item">
              <span className="graph-legend-dot" style={{ background: '#00ffcc' }} />
              Active + Main
            </div>
            <div className="graph-legend-item">
              <span className="graph-legend-dot" style={{ background: '#6c7aff' }} />
              Active branch
            </div>
            <div className="graph-legend-item">
              <span className="graph-legend-dot" style={{ background: '#00d4aa' }} />
              Main branch
            </div>
            <div className="graph-legend-item">
              <span className="graph-legend-dot" style={{ background: '#f0c040' }} />
              Generating
            </div>
            <div className="graph-legend-item">
              <span className="graph-legend-dot" style={{ background: '#555566' }} />
              Inactive
            </div>
            <div className="graph-legend-item">
              <span className="graph-legend-diamond" />
              Branch point
            </div>
          </div>
        )}
      </div>

      {contextMenu && (
        <div
          className="context-menu"
          style={{ left: contextMenu.x, top: contextMenu.y }}
          onClick={(e) => e.stopPropagation()}
        >
          <button className="context-menu-item" onClick={handleBranchFromGraph}>
            ⑂ Branch from here
          </button>
          <button className="context-menu-item" onClick={handleSetMainBranch}>
            ⭐ Set as Main Branch
          </button>
          <button
            className="context-menu-item"
            onClick={() => {
              handleNodeClick({ id: contextMenu.nodeId });
              setContextMenu(null);
            }}
          >
            👁 View this branch
          </button>
        </div>
      )}

      {branchNode && (
        <BranchDialog
          node={branchNode}
          onClose={() => setBranchNode(null)}
          onConfirm={(nodeId, newMessage) => {
            window.dispatchEvent(
              new CustomEvent('branchchat:branch', {
                detail: { nodeId, newMessage },
              })
            );
            setBranchNode(null);
          }}
        />
      )}
    </div>
  );
}

/**
 * Create a canvas-rendered sprite label for a graph node.
 * Uses THREE.Sprite + CanvasTexture instead of CSS2DObject to avoid
 * orphaned DOM elements (ghost text) when the graph re-renders.
 */
function createSpriteLabel(node, color) {
  const preview = node.userMessage
    ? node.userMessage.slice(0, 22) + (node.userMessage.length > 22 ? '…' : '')
    : 'Start';

  const turnNum = node.depth + 1;
  const statusIcon = node.status === 'generating' ? ' ⏳' : '';
  const turnText = `#${turnNum}${statusIcon}`;

  const canvas = document.createElement('canvas');
  const ctx = canvas.getContext('2d');

  // High-DPI canvas for crisp text
  const scale = 2;
  const canvasW = 200 * scale;
  const canvasH = 48 * scale;
  canvas.width = canvasW;
  canvas.height = canvasH;
  ctx.scale(scale, scale);

  // Turn number line
  ctx.font = 'bold 11px Inter, sans-serif';
  ctx.fillStyle = color;
  ctx.textAlign = 'center';
  ctx.textBaseline = 'top';
  ctx.fillText(turnText, 100, 4);

  // Preview line
  ctx.font = '9px Inter, sans-serif';
  ctx.fillStyle = 'rgba(232, 232, 240, 0.5)';
  ctx.fillText(preview, 100, 22);

  const texture = new THREE.CanvasTexture(canvas);
  texture.minFilter = THREE.LinearFilter;

  const spriteMat = new THREE.SpriteMaterial({
    map: texture,
    transparent: true,
    depthWrite: false,
  });
  const sprite = new THREE.Sprite(spriteMat);
  // Scale sprite to a reasonable world-space size
  sprite.scale.set(20, 20 * (canvasH / canvasW), 1);

  return sprite;
}

function escapeHtml(str) {
  return str
    .replace(/&/g, '&amp;')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;');
}