import { useEffect, useCallback, useState } from 'react';
import ReactFlow, { 
  Background, 
  Controls, 
  Node, 
  Edge, 
  useNodesState, 
  useEdgesState, 
  MarkerType,
  Position,
  Handle,
  NodeMouseHandler
} from 'reactflow';
import 'reactflow/dist/style.css';
// @ts-ignore
import ELK from 'elkjs/lib/elk.bundled.js';
import { useWorkflowStore } from '@/stores/useWorkflowStore';
import { schemas } from '@/api/schema.gen';
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 schemas.TaskStatus.enum.Completed: return <CheckCircle2 className="h-4 w-4 text-green-500" />;
        case schemas.TaskStatus.enum.Running: return <Loader2 className="h-4 w-4 text-blue-500 animate-spin" />;
        case schemas.TaskStatus.enum.Failed: return <AlertCircle className="h-4 w-4 text-red-500" />;
        case schemas.TaskStatus.enum.Skipped: return <Ban className="h-4 w-4 text-gray-400" />;
        case schemas.TaskStatus.enum.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, displayName?: string, status: TaskStatus, type: string }, selected: boolean }) => {
    const statusColors: Record<string, string> = {
        [schemas.TaskStatus.enum.Pending]: 'border-muted bg-card',
        [schemas.TaskStatus.enum.Scheduled]: 'border-yellow-500/50 bg-yellow-50/10',
        [schemas.TaskStatus.enum.Running]: 'border-blue-500 ring-2 ring-blue-500/20 bg-blue-50/10 animate-pulse',
        [schemas.TaskStatus.enum.Completed]: 'border-green-500 bg-green-50/10',
        [schemas.TaskStatus.enum.Failed]: 'border-red-500 bg-red-50/10',
        [schemas.TaskStatus.enum.Skipped]: 'border-gray-200 bg-gray-50/5 opacity-60',
    };

    // Use display name if available, otherwise format the label (ID/Name)
    const displayLabel = data.displayName || data.label.replace(/^(analysis:|fetch:)/, '').replace(/_/g, ' ');

    return (
        <div className={cn(
            "px-4 py-2 rounded-lg border shadow-sm min-w-[180px] transition-all duration-300 cursor-pointer bg-background",
            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 capitalize" title={displayLabel}>{displayLabel}</div>
            <Handle type="source" position={Position.Bottom} className="!bg-muted-foreground/50 !w-2 !h-2" />
        </div>
    );
};

const nodeTypes = {
    taskNode: WorkflowNode,
};

// --- Custom Grid Layout (Manual) ---
const useGridLayout = () => {
    const getLayoutedElements = useCallback(async (nodes: Node[], edges: Edge[], dagNodes: any[], dagEdges: any[]) => {
        // Constants
        const NODE_WIDTH = 240; // Increased width for long names
        const NODE_HEIGHT = 64;
        const X_SPACING = 40;
        const Y_SPACING = 80;
        const MAX_COLS = 3; // Max columns per row

        // 1. Topological Sort (Rank-Based Grid)
        // Calculate Ranks
        const ranks = new Map<string, number>();
        const nodeMap = new Map(dagNodes.map(n => [n.id, n]));
        const incomingEdges = new Map<string, string[]>();
        
        dagNodes.forEach(n => incomingEdges.set(n.id, []));
        dagEdges.forEach(e => incomingEdges.get(e.to)?.push(e.from));

        // Rank 0: Roots (No incoming edges)
        const roots = dagNodes.filter(n => (incomingEdges.get(n.id)?.length || 0) === 0);
        roots.forEach(n => ranks.set(n.id, 0));

        // Calculate Ranks for others (Longest Path)
        // Simple approach: For small DAGs, we can iterate. 
        // Correct approach: Kahn's algorithm or DFS with memoization
        
        const getRank = (nodeId: string, visited = new Set<string>()): number => {
            if (ranks.has(nodeId)) return ranks.get(nodeId)!;
            if (visited.has(nodeId)) return 0; // Cycle detected, fallback

            visited.add(nodeId);
            const parents = incomingEdges.get(nodeId) || [];
            if (parents.length === 0) return 0;

            let maxParentRank = 0;
            for (const pId of parents) {
                maxParentRank = Math.max(maxParentRank, getRank(pId, new Set(visited)));
            }
            
            const rank = maxParentRank + 1;
            ranks.set(nodeId, rank);
            return rank;
        };

        dagNodes.forEach(n => getRank(n.id));

        // Special Handling: 
        // 'company_profile' relies on DataFetch (Rank 0), so it gets Rank 1.
        // 'final_conclusion' relies on many nodes.
        // If 'final_conclusion' is the true sink, ensure it's at the bottom.
        // Let's identify the max rank.
        let maxRank = 0;
        ranks.forEach(r => maxRank = Math.max(maxRank, r));

        // Group nodes by Rank
        const nodesByRank: Record<number, any[]> = {};
        dagNodes.forEach(n => {
            const r = ranks.get(n.id) || 0;
            if (!nodesByRank[r]) nodesByRank[r] = [];
            nodesByRank[r].push(n);
        });

        // Sort nodes within ranks for consistency
        Object.values(nodesByRank).forEach(list => list.sort((a, b) => a.id.localeCompare(b.id)));

        // 2. Assign Grid Positions
        const layoutedNodes = [...nodes];
        let currentY = 0;

        // Iterate through ranks 0 to maxRank
        for (let r = 0; r <= maxRank; r++) {
            const rankNodes = nodesByRank[r] || [];
            if (rankNodes.length === 0) continue;

            // Grid layout for this rank
            // If it's Rank 0 (Root) or Max Rank (Sink), maybe center single row?
            // Actually, just applying grid logic uniformly works well.
            
            for (let i = 0; i < rankNodes.length; i += MAX_COLS) {
                const chunk = rankNodes.slice(i, i + MAX_COLS);
                const rowWidthTotal = chunk.length * NODE_WIDTH + (chunk.length - 1) * X_SPACING;
                const startX = -(rowWidthTotal / 2) + (NODE_WIDTH / 2);

                chunk.forEach((node, idx) => {
                    const n = layoutedNodes.find(ln => ln.id === node.id);
                    if (n) {
                        n.position = {
                            x: startX + idx * (NODE_WIDTH + X_SPACING),
                            y: currentY
                        };
                    }
                });
                
                currentY += (NODE_HEIGHT + Y_SPACING);
            }
        }

        // 3. Simplified Edges
        // Use default bezier curve for smoother look
        const layoutedEdges = edges.map(e => ({
            ...e,
            type: 'default', 
        }));

        return { nodes: layoutedNodes, edges: layoutedEdges, height: currentY };

    }, []);

    return { getLayoutedElements };
};

// --- Main Visualizer Component ---

export function WorkflowVisualizer() {
    const { dag, tasks, setActiveTab, activeTab } = useWorkflowStore();
    const [nodes, setNodes, onNodesChange] = useNodesState([]);
    const [edges, setEdges, onEdgesChange] = useEdgesState([]);
    const [graphHeight, setGraphHeight] = useState<number>(0);
    const { getLayoutedElements } = useGridLayout(); // Use custom layout

    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 createGraph = async () => {
            // 1. Create initial Nodes
            const initialNodes: Node[] = dag.nodes.map(node => {
                const taskState = tasks[node.id];
                const status = taskState ? taskState.status : node.initial_status;

                return {
                    id: node.id,
                    type: 'taskNode',
                    position: { x: 0, y: 0 },
                    draggable: false, // Disable node dragging
                    data: { 
                        label: node.name,
                        displayName: node.display_name,
                        status,
                        type: node.type
                    },
                    selected: activeTab === node.id,
                };
            });

            // 2. Create Edges
            const initialEdges: Edge[] = dag.edges.map(edge => {
                // Determine if this edge should be highlighted
                let isHighlighted = false;
                let isDimmed = false;

                if (activeTab === 'overview') {
                    // In overview, show all edges but subtly
                    isHighlighted = false;
                    isDimmed = false;
                } else {
                    // Specific node selected
                    if (edge.from === activeTab || edge.to === activeTab) {
                        isHighlighted = true;
                    } else {
                        isDimmed = true;
                    }
                }

                const style = isHighlighted 
                    ? { stroke: '#3b82f6', strokeWidth: 2, opacity: 1 } // Highlight: Blue, Thick, Opaque
                    : isDimmed 
                        ? { stroke: '#e2e8f0', strokeWidth: 1, opacity: 0.1 } // Dimmed: Faint
                        : { stroke: '#94a3b8', strokeWidth: 1, opacity: 0.4 }; // Normal (Overview): Medium gray

                return {
                    id: `${edge.from}-${edge.to}`,
                    source: edge.from,
                    target: edge.to,
                    markerEnd: { 
                        type: MarkerType.ArrowClosed,
                        color: style.stroke
                    },
                    animated: tasks[edge.from]?.status === schemas.TaskStatus.enum.Running || tasks[edge.to]?.status === schemas.TaskStatus.enum.Running,
                    zIndex: isHighlighted ? 10 : 0,
                    style
                };
            });

            // 3. Apply Layout
            // @ts-ignore - getLayoutedElements returns extra height prop
            const { nodes: layoutedNodes, edges: layoutedEdges, height } = await getLayoutedElements(initialNodes, initialEdges, dag.nodes, dag.edges);

            setNodes(layoutedNodes);
            setEdges(layoutedEdges);
            if (height) setGraphHeight(height);
        };

        createGraph();
    }, [dag, tasks, activeTab, setNodes, setEdges, getLayoutedElements]);

    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-full min-h-[300px] 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"
                panOnDrag={false} // Lock panning (X/Y) - user scroll does vertical
                zoomOnScroll={false} // Lock zoom
                zoomOnPinch={false}
                zoomOnDoubleClick={false}
                panOnScroll={true} // Allow vertical scroll to pan Y
                preventScrolling={false} // Allow browser scroll if needed, but usually false for canvas
                nodesConnectable={false} // Disable connecting nodes
                translateExtent={[
                    [-600, -100], 
                    [600, Math.max(graphHeight + 100, 500)] // Ensure at least some height
                ]}
            >
                <Background color="#94a3b8" gap={20} size={1} />
                {/* Controls removed to force simpler view */}
            </ReactFlow>
        </div>
    );
}
