frontend/src/components/MindMapModal.jsx

import React, { useState, useCallback, useEffect } from 'react';
import ReactFlow, {
    Background,
    Controls,
    MiniMap,
    useNodesState,
    useEdgesState,
    addEdge,
    MarkerType
} from 'reactflow';
import 'reactflow/dist/style.css';
import { X, Network, Maximize2, Sparkles } from 'lucide-react';

const MindMapModal = ({ isOpen, onClose, data, onExplainNode }) => {
    if (!isOpen) return null;

    const [nodes, setNodes, onNodesChange] = useNodesState([]);
    const [edges, setEdges, onEdgesChange] = useEdgesState([]);
    const [processingNode, setProcessingNode] = useState(null);

    // Initial Load
    useEffect(() => {
        if (!data) return;
        try {
            const parsedData = typeof data === 'string' ? JSON.parse(data) : data;

            // Root Node
            const rootNode = {
                id: 'root',
                type: 'default',
                data: {
                    label: parsedData.label || "Main Topic",
                    expandable: true,
                    expanded: true // Root is initially expanded if we load children 
                },
                position: { x: 0, y: 0 },
                style: {
                    background: '#fff',
                    border: '2px solid #3b82f6',
                    borderRadius: '8px',
                    padding: '12px',
                    width: 180,
                    fontSize: '14px',
                    fontWeight: '600',
                    color: '#1e293b',
                    boxShadow: '0 10px 15px -3px rgb(0 0 0 / 0.1)'
                }
            };

            const initialNodes = [rootNode];
            const initialEdges = [];

            // Parse initial children if any
            if (parsedData.children) {
                parsedData.children.forEach((child, i) => {
                    const childId = child.id || `child-${i}`;
                    initialNodes.push({
                        id: childId,
                        type: 'default',
                        data: {
                            label: child.label,
                            expandable: child.expandable !== false, // Default true unless specified
                            expanded: false,
                            has_children: child.has_children
                        },
                        position: { x: 300, y: (i - (parsedData.children.length - 1) / 2) * 150 },
                        style: {
                            background: '#fff',
                            border: '1px solid #e2e8f0',
                            borderRadius: '8px',
                            padding: '10px',
                            width: 160,
                            fontSize: '12px',
                            fontWeight: '500',
                            color: '#334155',
                            boxShadow: '0 4px 6px -1px rgb(0 0 0 / 0.1)'
                        }
                    });

                    initialEdges.push({
                        id: `root-${childId}`,
                        source: 'root',
                        target: childId,
                        type: 'smoothstep',
                        markerEnd: { type: MarkerType.ArrowClosed, color: '#94a3b8' },
                        style: { stroke: '#94a3b8', strokeWidth: 1.5 }
                    });
                });
            }

            setNodes(initialNodes);
            setEdges(initialEdges);

        } catch (e) {
            console.error("Failed to parse mind map data", e);
        }
    }, [data, setNodes, setEdges]);

    const onNodeClick = useCallback(async (event, node) => {
        // If already processing or already expanded, just focus or explain
        if (processingNode) return;

        const isLeaf = !node.data.expandable && !node.data.has_children;
        const isExpanded = node.data.expanded;

        // 1. LEAF NODE -> Explain
        if (isLeaf) {
            onExplainNode(node.data.label);
            return;
        }

        // 2. EXPANDABLE NODE -> Fetch Children
        if (!isExpanded && (node.data.expandable || node.data.has_children)) {
            setProcessingNode(node.id);

            // Visual feedback: change style to indicate loading
            setNodes(nds => nds.map(n => {
                if (n.id === node.id) {
                    return { ...n, style: { ...n.style, borderColor: '#f59e0b' }, data: { ...n.data, label: 'Loading...' } };
                }
                return n;
            }));

            try {
                // Dynamically import API here to avoid circular dependencies if any, or just use global
                const { expandMindMapNode } = await import('../api');
                const result = await expandMindMapNode(node.data.label);

                // Parse result
                const childrenData = typeof result.answer === 'string' ? JSON.parse(result.answer).children : result.answer.children;

                if (!childrenData || childrenData.length === 0) {
                    // No children found, treat as leaf
                    onExplainNode(node.data.label);
                    setNodes(nds => nds.map(n => n.id === node.id ? { ...n, data: { ...n.data, label: node.data.label, expandable: false } } : n));
                    return;
                }

                // Add new nodes
                const newNodes = [];
                const newEdges = [];
                const parentX = node.position.x;
                const parentY = node.position.y;

                childrenData.forEach((child, i) => {
                    const childId = child.id || `${node.id}-child-${i}-${Math.random().toString(36).substr(2, 9)}`;

                    newNodes.push({
                        id: childId,
                        type: 'default',
                        data: {
                            label: child.label,
                            expandable: child.has_children !== false,
                            expanded: false,
                            has_children: child.has_children
                        },
                        // Position relative to parent
                        position: {
                            x: parentX + 300,
                            y: parentY + (i - (childrenData.length - 1) / 2) * 120
                        },
                        style: {
                            background: '#fff',
                            border: '1px solid #e2e8f0',
                            borderRadius: '8px',
                            padding: '10px',
                            width: 160,
                            fontSize: '12px',
                            fontWeight: '500',
                            color: '#334155',
                            boxShadow: '0 4px 6px -1px rgb(0 0 0 / 0.1)'
                        }
                    });

                    newEdges.push({
                        id: `${node.id}-${childId}`,
                        source: node.id,
                        target: childId,
                        type: 'smoothstep',
                        markerEnd: { type: MarkerType.ArrowClosed, color: '#94a3b8' },
                        style: { stroke: '#94a3b8', strokeWidth: 1.5 }
                    });
                });

                // Update state
                setNodes(nds => nds.map(n => {
                    if (n.id === node.id) {
                        return {
                            ...n,
                            style: { ...n.style, borderColor: '#3b82f6' }, // Reset color
                            data: { ...n.data, label: node.data.label, expanded: true }
                        };
                    }
                    return n;
                }).concat(newNodes));

                setEdges(eds => eds.concat(newEdges));

            } catch (error) {
                console.error("Error expanding node:", error);
                // Reset state on error
                setNodes(nds => nds.map(n => n.id === node.id ? { ...n, style: { ...n.style, borderColor: '#ef4444' }, data: { ...n.data, label: node.data.label } } : n));
            } finally {
                setProcessingNode(null);
            }
        } else {
            // Already expanded, maybe explain context too or collapse (collapse logic omitted for simplicity/Task 7 rules says click expand OR explain)
            // If already expanded, let's explain it
            onExplainNode(node.data.label);
        }

    }, [nodes, processingNode, setNodes, setEdges, onExplainNode]);

    return (
        <div className="fixed inset-0 z-50 flex items-center justify-center bg-slate-900/80 backdrop-blur-sm p-4 animate-in fade-in duration-200">
            <div className="bg-white w-full h-full rounded-2xl shadow-2xl flex flex-col overflow-hidden relative border border-white/20">
                {/* Header */}
                <div className="bg-white border-b border-slate-100 p-4 flex justify-between items-center z-10 shadow-sm shrink-0">
                    <div className="flex items-center gap-2 text-slate-700">
                        <div className="p-2 bg-blue-50 text-blue-600 rounded-lg">
                            <Network size={20} />
                        </div>
                        <div>
                            <h2 className="font-bold text-lg leading-tight">Interactive Mind Map</h2>
                            <p className="text-xs text-slate-500">
                                {processingNode ? "Expanding concept..." : "Click to expand • Leaf nodes explain concept"}
                            </p>
                        </div>
                    </div>
                    <button onClick={onClose} className="p-2 hover:bg-slate-100 rounded-full text-slate-500 transition-colors">
                        <X size={24} />
                    </button>
                </div>

                {/* Canvas */}
                <div className="flex-1 bg-slate-50 relative">
                    <ReactFlow
                        nodes={nodes}
                        edges={edges}
                        onNodesChange={onNodesChange}
                        onEdgesChange={onEdgesChange}
                        onNodeClick={onNodeClick}
                        fitView
                        attributionPosition="bottom-right"
                        minZoom={0.1}
                        maxZoom={2}
                    >
                        <Background color="#cbd5e1" gap={20} size={1} />
                        <Controls className="bg-white border border-slate-200 shadow-sm text-slate-600" />
                        <MiniMap
                            className="border border-slate-200 shadow-sm rounded-lg overflow-hidden"
                            nodeColor={n => n.type === 'input' ? '#3b82f6' : '#fff'}
                        />
                    </ReactFlow>

                    {/* Loading Overlay if needed, or just rely on node state */}
                </div>
            </div>
        </div>
    );
};

export default MindMapModal;