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

const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
  const { updateNodeData } = useFlowStore();
  const updateNodeInternals = useUpdateNodeInternals();
  const edges = useEdges();

  // Force update handles when traces change
  useEffect(() => {
    updateNodeInternals(id);
  }, [id, data.outgoingTraces, 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);

  return (
    <div className={`px-4 py-2 shadow-md rounded-md bg-white border-2 min-w-[200px] ${selected ? 'border-blue-500' : 'border-gray-200'}`}>
      <div className="flex items-center mb-2">
        <div className="rounded-full w-8 h-8 flex justify-center items-center bg-gray-100">
          {data.status === 'loading' ? (
            <Loader2 className="w-4 h-4 animate-spin text-blue-500" />
          ) : (
            <MessageSquare className="w-4 h-4 text-gray-600" />
          )}
        </div>
        <div className="ml-2">
          <div className="text-sm font-bold truncate max-w-[150px]">{data.label}</div>
          <div className="text-xs 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">
         {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;
           
           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]"
                  style={{ 
                    top: '50%', 
                    transform: 'translateY(-50%)',
                    backgroundColor: edgeColor || '#3b82f6', // Default blue if not connected
                    border: edgeColor ? 'none' : undefined 
                  }}
               />
               <span className="absolute left-4 top-[-2px] text-[9px] text-gray-400 pointer-events-none">
                 {i}
               </span>
             </div>
           );
         })}
      </div>

      {/* Dynamic Outputs (Traces) */}
      <div className="absolute right-0 top-0 bottom-0 flex flex-col justify-center w-4">
          {/* 1. Outgoing Traces (Pass-through + Self) */}
          {data.outgoingTraces && data.outgoingTraces.map((trace, i) => (
             <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={{ 
                     backgroundColor: trace.color,
                     top: '50%', 
                     transform: 'translateY(-50%)' 
                   }}
                />
             </div>
          ))}
          
          {/* 2. 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] text-gray-400 pointer-events-none w-max">
               + New
             </span>
          </div>
      </div>
      
      {data.status === 'error' && (
        <div className="text-xs text-red-500 mt-2">Error</div>
      )}
    </div>
  );
};

export default LLMNode;