path: root/frontend/src/components/nodes/LLMNode.tsx

import { useEffect, useState } from 'react';
import { Handle, Position, type NodeProps, useUpdateNodeInternals, useEdges } from 'reactflow';
import type { NodeData, MergedTrace } from '../../store/flowStore';
import { Loader2, MessageSquare } from 'lucide-react';
import useFlowStore from '../../store/flowStore';

const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
  const { theme, nodes } = useFlowStore();
  const [showPreview, setShowPreview] = useState(false);
  const updateNodeInternals = useUpdateNodeInternals();
  const edges = useEdges();

  // Force update handles when traces change
  useEffect(() => {
    updateNodeInternals(id);
  }, [id, data.outgoingTraces, data.mergedTraces, data.inputs, updateNodeInternals]);

  // Determine how many input handles to show
  // We want to ensure there is always at least one empty handle at the bottom
  // plus all currently connected handles.
  
  // Find all edges connected to this node's inputs
  const connectedHandles = new Set(
    edges
      .filter(e => e.target === id)
      .map(e => e.targetHandle)
  );
  
  // Logic: 
  // If input-0 is connected, show input-1.
  // If input-1 is connected, show input-2.
  // We can just iterate until we find an unconnected one.
  
  let handleCount = 1;
  while (connectedHandles.has(`input-${handleCount - 1}`)) {
    handleCount++;
  }
  
  // But wait, if we delete an edge to input-0, we still want input-1 to exist if it's connected?
  // No, usually in this designs, we just render up to max(connected_index) + 1.
  
  // Let's get the max index connected
  let maxConnectedIndex = -1;
  edges.filter(e => e.target === id).forEach(e => {
    const idx = parseInt(e.targetHandle?.replace('input-', '') || '0');
    if (!isNaN(idx) && idx > maxConnectedIndex) {
      maxConnectedIndex = idx;
    }
  });
  
  const inputsToShow = Math.max(maxConnectedIndex + 2, 1);

  const isDisabled = data.disabled;
  const isDark = theme === 'dark';

  // Truncate preview content
  const previewContent = data.response 
    ? data.response.slice(0, 200) + (data.response.length > 200 ? '...' : '')
    : data.userPrompt 
      ? data.userPrompt.slice(0, 100) + (data.userPrompt.length > 100 ? '...' : '')
      : null;

  return (
    <div 
      className={`px-4 py-2 shadow-md rounded-md border-2 min-w-[200px] transition-all relative ${
        isDisabled 
          ? isDark 
            ? 'bg-gray-800 border-gray-600 opacity-50 cursor-not-allowed' 
            : 'bg-gray-100 border-gray-300 opacity-50 cursor-not-allowed'
          : selected 
            ? isDark
              ? 'bg-gray-800 border-blue-400'
              : 'bg-white border-blue-500'
            : isDark
              ? 'bg-gray-800 border-gray-600'
              : 'bg-white border-gray-200'
      }`}
      style={{ pointerEvents: isDisabled ? 'none' : 'auto' }}
      onMouseEnter={() => setShowPreview(true)}
      onMouseLeave={() => setShowPreview(false)}
    >
      {/* Content Preview Tooltip */}
      {showPreview && previewContent && !isDisabled && (
        <div 
          className={`absolute z-50 left-1/2 -translate-x-1/2 bottom-full mb-2 w-64 p-3 rounded-lg shadow-xl text-xs whitespace-pre-wrap pointer-events-none ${
            isDark ? 'bg-gray-700 text-gray-200 border border-gray-600' : 'bg-white text-gray-700 border border-gray-200'
          }`}
        >
          <div className={`font-semibold mb-1 ${isDark ? 'text-gray-400' : 'text-gray-500'}`}>
            {data.response ? 'Response Preview' : 'Prompt Preview'}
          </div>
          {previewContent}
          {/* Arrow */}
          <div className={`absolute left-1/2 -translate-x-1/2 top-full w-0 h-0 border-l-8 border-r-8 border-t-8 border-l-transparent border-r-transparent ${
            isDark ? 'border-t-gray-700' : 'border-t-white'
          }`} />
        </div>
      )}

      <div className="flex items-center mb-2">
        <div className={`rounded-full w-8 h-8 flex justify-center items-center ${
          isDisabled 
            ? isDark ? 'bg-gray-700' : 'bg-gray-200' 
            : isDark ? 'bg-gray-700' : 'bg-gray-100'
        }`}>
          {data.status === 'loading' ? (
            <Loader2 className="w-4 h-4 animate-spin text-blue-500" />
          ) : (
            <MessageSquare className={`w-4 h-4 ${
              isDisabled 
                ? 'text-gray-500' 
                : isDark ? 'text-gray-400' : 'text-gray-600'
            }`} />
          )}
        </div>
        <div className="ml-2">
          <div className={`text-sm font-bold truncate max-w-[150px] ${
            isDisabled 
              ? 'text-gray-500' 
              : isDark ? 'text-gray-200' : 'text-gray-900'
          }`}>
            {data.label}
            {isDisabled && <span className="text-xs ml-1">(disabled)</span>}
          </div>
          <div className={`text-xs ${isDark ? 'text-gray-400' : 'text-gray-500'}`}>{data.model}</div>
        </div>
      </div>

      {/* Dynamic Inputs */}
      <div className="absolute left-0 top-0 bottom-0 flex flex-col justify-center w-4">
         {/* Regular input handles */}
         {Array.from({ length: inputsToShow }).map((_, i) => {
           // Find the connected edge to get color
           const connectedEdge = edges.find(e => e.target === id && e.targetHandle === `input-${i}`);
           const edgeColor = connectedEdge?.style?.stroke as string;
           
           // Check if this is a merged trace connection
           const edgeData = (connectedEdge as any)?.data;
           const isMergedTrace = edgeData?.isMerged;
           const mergedColors = edgeData?.colors as string[] | undefined;
           
           // Create gradient for merged traces
           let handleBackground: string = edgeColor || '#3b82f6';
           if (isMergedTrace && mergedColors && mergedColors.length >= 2) {
             const gradientStops = mergedColors.map((color, idx) => 
               `${color} ${(idx / mergedColors.length) * 100}%, ${color} ${((idx + 1) / mergedColors.length) * 100}%`
             ).join(', ');
             handleBackground = `linear-gradient(45deg, ${gradientStops})`;
           }
           
           return (
             <div key={i} className="relative h-4 w-4 my-1">
               <Handle 
                  type="target" 
                  position={Position.Left} 
                  id={`input-${i}`}
                  className="!w-3 !h-3 !left-[-6px] !border-0"
                  style={{ 
                    top: '50%', 
                    transform: 'translateY(-50%)',
                    background: handleBackground
                  }}
               />
               <span className={`absolute left-4 top-[-2px] text-[9px] pointer-events-none ${isDark ? 'text-gray-500' : 'text-gray-400'}`}>
                 {i}
               </span>
             </div>
           );
         })}
         
         {/* Prepend input handles for traces that this node is the HEAD of */}
         {/* Show dashed handle if no prepend connection yet (can accept prepend) */}
         {/* Show solid handle if already has prepend connection (connected) */}
         {data.outgoingTraces && data.outgoingTraces
           .filter(trace => {
             // Check if this is a self trace, fork trace originated from this node
             const isSelfTrace = trace.id === `trace-${id}`;
             // Strict check for fork trace: must be originated from this node
             const isForkTrace = trace.id.startsWith('fork-') && trace.sourceNodeId === id;
             
             if (!isSelfTrace && !isForkTrace) return false;
             
             // Check if this trace has any outgoing edges (downstream connections)
             const hasDownstream = edges.some(e => 
               e.source === id && e.sourceHandle === `trace-${trace.id}`
             );
             return hasDownstream;
           })
           .map((trace) => {
             // Check if there's already a prepend connection to this trace
             const hasPrependConnection = edges.some(e => 
               e.target === id && e.targetHandle === `prepend-${trace.id}`
             );
             const prependEdge = edges.find(e => 
               e.target === id && e.targetHandle === `prepend-${trace.id}`
             );
             
             return (
               <div key={`prepend-${trace.id}`} className="relative h-4 w-4 my-1" title={`Prepend context to: ${trace.id}`}>
                 <Handle 
                    type="target" 
                    position={Position.Left} 
                    id={`prepend-${trace.id}`}
                    className="!w-3 !h-3 !left-[-6px] !border-2"
                    style={{ 
                      top: '50%', 
                      transform: 'translateY(-50%)',
                      backgroundColor: hasPrependConnection 
                        ? (prependEdge?.style?.stroke as string || trace.color)
                        : trace.color,
                      borderColor: isDark ? '#374151' : '#fff',
                      borderStyle: hasPrependConnection ? 'solid' : 'dashed'
                    }}
                 />
               </div>
             );
           })
         }
         
         {/* Prepend input handles for merged traces - REMOVED per user request */}
         {/* Users should prepend to parent traces instead */}
      </div>

      {/* Dynamic Outputs (Traces) */}
      <div className="absolute right-0 top-0 bottom-0 flex flex-col justify-center w-4">
          {/* 0. Incoming Trace Continue Handles - allow continuing an incoming trace */}
          {/* Only show if there's NO downstream edge yet (dashed = waiting for connection) */}
          {data.traces && data.traces
            .filter((trace: Trace) => {
              // Only show continue handle if NOT already connected downstream
              // Now that trace IDs don't evolve, we check the base ID
              const hasOutgoingEdge = edges.some(e => 
                e.source === id && 
                e.sourceHandle === `trace-${trace.id}`
              );
              // Only show dashed handle if not yet connected
              return !hasOutgoingEdge;
            })
            .map((trace: Trace) => {
              // Handle merged trace visualization
              let backgroundStyle = trace.color;
              if (trace.isMerged && trace.mergedColors && trace.mergedColors.length > 0) {
                const colors = trace.mergedColors;
                const gradientStops = colors.map((color, idx) => 
                  `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
                ).join(', ');
                backgroundStyle = `linear-gradient(45deg, ${gradientStops})`;
              }

              return (
                <div key={`continue-${trace.id}`} className="relative h-4 w-4 my-1" title={`Continue trace: ${trace.id}`}>
                  <Handle 
                     type="source" 
                     position={Position.Right} 
                     id={`trace-${trace.id}`}
                     className="!w-3 !h-3 !right-[-6px]"
                     style={{ 
                       background: backgroundStyle,
                       top: '50%', 
                       transform: 'translateY(-50%)',
                       border: `2px dashed ${isDark ? '#374151' : '#fff'}`
                     }}
                  />
                </div>
              );
            })}
          
          {/* 1. Regular Outgoing Traces (Self + Forks + Connected Pass-through) */}
          {/* Only show traces that have actual downstream edges */}
          {data.outgoingTraces && data.outgoingTraces
            .filter(trace => {
              // Check if this is a locally created merged trace (should be shown in Part 2)
              const isLocallyMerged = data.mergedTraces?.some(m => m.id === trace.id);
              if (isLocallyMerged) return false;
              
              // Only show if there's an actual downstream edge using this trace
              const hasDownstream = edges.some(e => 
                e.source === id && e.sourceHandle === `trace-${trace.id}`
              );
              return hasDownstream;
            })
            .map((trace) => {
             // Handle merged trace visualization
             let backgroundStyle = trace.color;
             if (trace.isMerged && trace.mergedColors && trace.mergedColors.length > 0) {
               const colors = trace.mergedColors;
               const gradientStops = colors.map((color, idx) => 
                 `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
               ).join(', ');
               backgroundStyle = `linear-gradient(45deg, ${gradientStops})`;
             }

             return (
             <div key={trace.id} className="relative h-4 w-4 my-1" title={`Trace: ${trace.id}`}>
                <Handle 
                   type="source" 
                   position={Position.Right} 
                   id={`trace-${trace.id}`}
                   className="!w-3 !h-3 !right-[-6px]"
                   style={{ 
                     background: backgroundStyle,
                     top: '50%', 
                     transform: 'translateY(-50%)' 
                   }}
                />
             </div>
             );
          })}
          
          {/* 2. Merged Trace Handles (with alternating color stripes) */}
          {data.mergedTraces && data.mergedTraces.map((merged: MergedTrace) => {
            // Check if this merged trace has any outgoing edges
            const hasOutgoingEdge = edges.some(e => 
              e.source === id && e.sourceHandle === `trace-${merged.id}`
            );
            
            // Create a gradient background from the source trace colors
            const colors = merged.colors.length > 0 ? merged.colors : ['#888'];
            const gradientStops = colors.map((color, idx) => 
              `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
            ).join(', ');
            const stripeGradient = `linear-gradient(45deg, ${gradientStops})`;
            
            return (
              <div key={merged.id} className="relative h-4 w-4 my-1" title={`Merged: ${merged.strategy} (${merged.sourceTraceIds.length} traces)`}>
                <Handle 
                   type="source" 
                   position={Position.Right} 
                   id={`trace-${merged.id}`}
                   className="!w-3 !h-3 !right-[-6px]"
                   style={{ 
                     background: stripeGradient,
                     top: '50%', 
                     transform: 'translateY(-50%)',
                     border: hasOutgoingEdge ? 'none' : `2px dashed ${isDark ? '#374151' : '#fff'}`
                   }}
                />
              </div>
            );
          })}
          
          {/* 3. New Branch Generator Handle (Always visible) */}
          <div className="relative h-4 w-4 my-1" title="Create New Branch">
             <Handle 
                type="source" 
                position={Position.Right} 
                id="new-trace"
                className="!w-3 !h-3 !bg-gray-400 !right-[-6px]"
                style={{ top: '50%', transform: 'translateY(-50%)' }}
             />
             <span className={`absolute right-4 top-[-2px] text-[9px] pointer-events-none w-max ${isDark ? 'text-gray-500' : 'text-gray-400'}`}>
               + New
             </span>
          </div>
      </div>
      
      {data.status === 'error' && (
        <div className="text-xs text-red-500 mt-2">Error</div>
      )}
    </div>
  );
};

export default LLMNode;