import {Background, Controls, MiniMap, ReactFlow, useEdgesState, useNodesState, Position, Handle, ReactFlowProvider} from '@xyflow/react';
import React, {useCallback, memo, useMemo, useRef, useState, useEffect, forwardRef, useImperativeHandle} from 'react';
import {Play} from 'lucide-react';
import {useLoadTest} from '@/providers/LoadTestProvider.jsx';

// Optimized node components with proper prop types and memoization
const StartNode = memo(({data, isConnectable}) => {
  // Memoize handles configuration
  const handles = useMemo(() => ([
    {id: 'start-right', type: 'source', position: Position.Right}
  ]), []);

  return (
      <div className="bg-green-50 rounded-lg border-2 border-green-500 shadow-lg p-4 w-[100px]">
        {handles.map(handle => (
            <Handle
                key={handle.id}
                type={handle.type}
                position={handle.position}
                isConnectable={isConnectable}
                className="w-3 h-3 bg-green-500"
            />
        ))}
        <div className="flex items-center justify-center gap-2">
          <Play className="h-4 w-4 text-green-600"/>
          <span className="font-semibold text-sm text-green-700">Start</span>
        </div>
      </div>
  );
});

const ApiNode = memo(({data, isConnectable}) => {
  // Memoize handles and colors
  const handles = useMemo(() => ([
    {id: `${data.id}-left`, type: 'target', position: Position.Left},
    {id: `${data.id}-right`, type: 'source', position: Position.Right}
  ]), [data.id]);

  const methodColor = useMemo(() => {
    const colors = {
      GET: 'bg-green-500',
      POST: 'bg-blue-500',
      PUT: 'bg-yellow-500',
      DELETE: 'bg-red-500',
      default: 'bg-gray-500'
    };
    return colors[data.method] || colors.default;
  }, [data.method]);

  return (
      <div className="bg-white rounded-lg border shadow-lg p-4 w-[200px]">
        {handles.map(handle => (
            <Handle
                key={handle.id}
                type={handle.type}
                position={handle.position}
                isConnectable={isConnectable}
                className="w-3 h-3 bg-gray-500"
            />
        ))}
        <div className="flex items-center justify-between mb-2">
          <div className="flex items-center gap-2">
            <div className={`w-2 h-2 rounded-full ${methodColor}`}/>
            <span className="font-bold text-sm">{data.method}</span>
          </div>
        </div>
        <div className="text-xs text-gray-500">{data.name}</div>
        <div className="text-sm truncate text-gray-600 mb-2">
          {data.path}
        </div>
      </div>
  );
});

const ConditionNode = memo(({data, isConnectable}) => {
  // Memoize handle styles
  const handleStyles = useMemo(() => ({
    true: {
      top: '30%',
      backgroundColor: '#22c55e',
      border: '2px solid #16a34a'
    },
    false: {
      top: '70%',
      backgroundColor: '#ef4444',
      border: '2px solid #dc2626'
    }
  }), []);

  return (
      <div className="bg-white rounded-lg border shadow-lg p-4 w-[200px]">
        <Handle
            type="target"
            position={Position.Left}
            id="in"
            isConnectable={isConnectable}
            className="w-3 h-3"
        />
        <Handle
            type="source"
            position={Position.Right}
            id="true"
            style={handleStyles.true}
            isConnectable={isConnectable}
            className="w-3 h-3 !opacity-100"
        />
        <Handle
            type="source"
            position={Position.Right}
            id="false"
            style={handleStyles.false}
            isConnectable={isConnectable}
            className="w-3 h-3 !opacity-100"
        />
        <div className="flex items-center justify-between mb-2">
          <div className="flex items-center gap-2">
            <div className="w-2 h-2 rounded-full bg-purple-500"/>
            <span className="font-bold text-sm">Condition</span>
          </div>
        </div>
        <div className="text-sm text-gray-600">
          {data.condition || 'Add condition'}
        </div>
      </div>
  );
});

const DelayNode = memo(({data, isConnectable}) => {
  const handles = useMemo(() => ([
    {id: `${data.id}-left`, type: 'target', position: Position.Left},
    {id: `${data.id}-right`, type: 'source', position: Position.Right}
  ]), [data.id]);

  return (
      <div className="bg-white rounded-lg border shadow-lg p-4 w-[200px]">
        {handles.map(handle => (
            <Handle
                key={handle.id}
                type={handle.type}
                position={handle.position}
                isConnectable={isConnectable}
                className="w-3 h-3 bg-gray-500"
            />
        ))}
        <div className="flex items-center justify-between mb-2">
          <div className="flex items-center gap-2">
            <div className="w-2 h-2 rounded-full bg-orange-500"/>
            <span className="font-bold text-sm">Delay</span>
          </div>
        </div>
        <div className="text-sm text-gray-600">
          Delay: {data.duration || '1000'}ms
        </div>
      </div>
  );
});

// Add display names for debugging
StartNode.displayName = 'StartNode';
ApiNode.displayName = 'ApiNode';
ConditionNode.displayName = 'ConditionNode';
DelayNode.displayName = 'DelayNode';

// Memoize node types
const nodeTypes = {
  startNode: StartNode,
  apiNode: ApiNode,
  conditionNode: ConditionNode,
  delayNode: DelayNode
};

const DiagramArea = forwardRef(({onSave, onFlowChange, isExecuting}, ref) => {
  const {currentFlowData} = useLoadTest();
  const reactFlowWrapper = useRef(null);
  const [selectedNodeId, setSelectedNodeId] = useState(null);
  const [edges, setEdges, onEdgesChange] = useEdgesState(currentFlowData?.edges || []);
  const [nodes, setNodes, onNodesChange] = useNodesState(currentFlowData?.nodes || [
    {
      id: 'start',
      type: 'startNode',
      position: {x: 100, y: 50},
      data: {id: 'start'}
    }]);

  useEffect(() => {
    if (currentFlowData) {
      const {nodes: scenarioNodes, edges: scenarioEdges} = currentFlowData;
      setNodes(scenarioNodes?.length ? scenarioNodes : [
        {
          id: 'start',
          type: 'startNode',
          position: {x: 100, y: 50},
          data: {id: 'start'}
        }]);
      setEdges(scenarioEdges || []);
    }
  }, [currentFlowData, setNodes, setEdges]);

  useEffect(() => {
    onFlowChange?.(nodes, edges);
  }, [nodes, edges, onFlowChange]);

  useImperativeHandle(ref, () => ({
    getFlowData: () => ({
      nodes,
      edges,
      viewport: reactFlowWrapper.current?.getViewport()
    })
  }));

  // Click handlers for node selection
  const onNodeClick = useCallback((event, node) => {
    event.stopPropagation();
    setSelectedNodeId(node.id);
  }, []);

  const onPaneClick = useCallback(() => {
    setSelectedNodeId(null);
  }, []);

  // Apply styles to nodes
  useEffect(() => {
    setNodes((nds) =>
        nds.map((node) => ({
          ...node,
          className: node.id === selectedNodeId ? 'border-2 border-primary shadow-md bg-primary/5' : null
        }))
    );
  }, [selectedNodeId, setNodes]);

  const onConnect = useCallback((params) => {
    setEdges(prev => {
      const newId = `edge-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
      return [...prev, {...params, id: newId}];
    });
  }, []);

  const onInit = useCallback((instance) => {
    reactFlowWrapper.current = instance;
  }, []);

  const onDragOver = useCallback((event) => {
    event.preventDefault();
    event.dataTransfer.dropEffect = 'copy';
  }, []);

  const onDrop = useCallback((event) => {
    event.preventDefault();

    if (!reactFlowWrapper.current) {
      return;
    }

    try {
      const apiData = JSON.parse(event.dataTransfer.getData('application/json'));
      const position = reactFlowWrapper.current.screenToFlowPosition({
        x: event.clientX,
        y: event.clientY
      });

      const newNode = {
        id: `node-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
        type: 'apiNode',
        position,
        data: {
          ...apiData,
          id: `api-${Date.now()}` // Ensure unique API ID
        }
      };

      setNodes(prev => prev.concat(newNode));
    } catch (error) {
      console.error('Error adding new node:', error);
    }
  }, [setNodes]);

  // Default edge options memoized
  const defaultEdgeOptions = useMemo(() => ({type: 'smoothstep'}), []);

  return (

      <ReactFlowProvider>
        <ReactFlow
            nodes={nodes}
            edges={edges}
            onNodesChange={onNodesChange}
            onEdgesChange={onEdgesChange}
            onConnect={onConnect}
            onInit={onInit}
            defaultEdgeOptions={defaultEdgeOptions}
            onDrop={onDrop}
            onDragOver={onDragOver}
            onNodeClick={onNodeClick}
            onPaneClick={onPaneClick}
            nodeTypes={nodeTypes}
            proOptions={{hideAttribution: true}}
        >
          <Background/>
          <Controls/>
          <MiniMap/>
        </ReactFlow>
      </ReactFlowProvider>

  );
});

DiagramArea.displayName = 'DiagramArea';

export default memo(DiagramArea);
