import React, { useEffect, useCallback } from 'react';
import ReactFlow, { 
  Background, 
  Controls, 
  Node, 
  Edge, 
  useNodesState, 
  useEdgesState, 
  MarkerType,
  Position,
  Handle,
  useOnSelectionChange,
  NodeMouseHandler
} from 'reactflow';
import 'reactflow/dist/style.css';
import { useWorkflowStore } from '@/stores/useWorkflowStore';
import { TaskStatus } from '@/types/workflow';
import { cn } from '@/lib/utils';
import { CheckCircle2, Circle, Clock, AlertCircle, Loader2, Ban } from 'lucide-react';

// --- Custom Node Component ---
const StatusIcon = ({ status }: { status: TaskStatus }) => {
    switch (status) {
        case 'Completed': return <CheckCircle2 className="h-4 w-4 text-green-500" />;
        case 'Running': return <Loader2 className="h-4 w-4 text-blue-500 animate-spin" />;
        case 'Failed': return <AlertCircle className="h-4 w-4 text-red-500" />;
        case 'Skipped': return <Ban className="h-4 w-4 text-gray-400" />;
        case 'Scheduled': return <Clock className="h-4 w-4 text-yellow-500" />;
        default: return <Circle className="h-4 w-4 text-muted-foreground" />;
    }
};

const WorkflowNode = ({ data, selected }: { data: { label: string, status: TaskStatus, type: string }, selected: boolean }) => {
    const statusColors = {
        'Pending': 'border-muted bg-card',
        'Scheduled': 'border-yellow-500/50 bg-yellow-50/10',
        'Running': 'border-blue-500 ring-2 ring-blue-500/20 bg-blue-50/10',
        'Completed': 'border-green-500 bg-green-50/10',
        'Failed': 'border-red-500 bg-red-50/10',
        'Skipped': 'border-gray-200 bg-gray-50/5 opacity-60',
    };

    return (
        <div className={cn(
            "px-4 py-2 rounded-lg border shadow-sm min-w-[150px] transition-all duration-300 cursor-pointer",
            statusColors[data.status] || 'border-muted',
            selected && "ring-2 ring-primary border-primary"
        )}>
            <Handle type="target" position={Position.Top} className="!bg-muted-foreground/50 !w-2 !h-2" />
            <div className="flex items-center gap-2 mb-1">
                <StatusIcon status={data.status} />
                <div className="text-xs font-mono uppercase text-muted-foreground">{data.type}</div>
            </div>
            <div className="font-medium text-sm truncate" title={data.label}>{data.label}</div>
            <Handle type="source" position={Position.Bottom} className="!bg-muted-foreground/50 !w-2 !h-2" />
        </div>
    );
};

const nodeTypes = {
    taskNode: WorkflowNode,
};

// --- Main Visualizer Component ---

export function WorkflowVisualizer() {
    const { dag, tasks, setActiveTab, activeTab } = useWorkflowStore();
    const [nodes, setNodes, onNodesChange] = useNodesState([]);
    const [edges, setEdges, onEdgesChange] = useEdgesState([]);

    const onNodeClick: NodeMouseHandler = useCallback((event, node) => {
        setActiveTab(node.id);
    }, [setActiveTab]);

    // Transform DAG to ReactFlow nodes/edges when DAG or Task Status changes
    useEffect(() => {
        if (!dag) return;

        const levels: Record<string, number> = {};
        const getLevel = (id: string): number => {
            if (levels[id] !== undefined) return levels[id];
            const node = dag.nodes.find(n => n.id === id);
            if (!node || node.dependencies.length === 0) {
                levels[id] = 0;
                return 0;
            }
            const maxParentLevel = Math.max(...node.dependencies.map(getLevel));
            levels[id] = maxParentLevel + 1;
            return maxParentLevel + 1;
        };

        dag.nodes.forEach(n => getLevel(n.id));

        const levelCounts: Record<number, number> = {};
        const newNodes: Node[] = dag.nodes.map(node => {
            const level = levels[node.id];
            const indexInLevel = levelCounts[level] || 0;
            levelCounts[level] = indexInLevel + 1;

            const taskState = tasks[node.id];
            const status = taskState ? taskState.status : node.initial_status;

            return {
                id: node.id,
                type: 'taskNode',
                position: { x: indexInLevel * 220, y: level * 120 }, 
                data: { 
                    label: node.label || node.id,
                    status,
                    type: node.type
                },
                selected: activeTab === node.id, // Highlight active node
                targetPosition: Position.Top,
                sourcePosition: Position.Bottom,
            };
        });

        const newEdges: Edge[] = [];
        dag.nodes.forEach(node => {
            node.dependencies.forEach(depId => {
                newEdges.push({
                    id: `${depId}-${node.id}`,
                    source: depId,
                    target: node.id,
                    type: 'smoothstep',
                    markerEnd: { type: MarkerType.ArrowClosed },
                    animated: tasks[depId]?.status === 'Running' || tasks[node.id]?.status === 'Running',
                    style: { stroke: '#64748b', strokeWidth: 1.5 }
                });
            });
        });

        setNodes(newNodes);
        setEdges(newEdges);
    }, [dag, tasks, activeTab, setNodes, setEdges]); // Depend on activeTab for selection

    if (!dag) return <div className="flex items-center justify-center h-full text-muted-foreground">Waiting for workflow to start...</div>;

    return (
        <div className="h-[400px] w-full border rounded-lg bg-muted/5">
            <ReactFlow
                nodes={nodes}
                edges={edges}
                onNodesChange={onNodesChange}
                onEdgesChange={onEdgesChange}
                onNodeClick={onNodeClick}
                nodeTypes={nodeTypes}
                fitView
                attributionPosition="bottom-right"
            >
                <Background color="#94a3b8" gap={20} size={1} />
                <Controls />
            </ReactFlow>
        </div>
    );
}
