import React, { useState, useCallback } from 'react';
import {
    ReactFlow,
    Controls,
    Background,
    applyNodeChanges,
    applyEdgeChanges,
    addEdge,
    Node,
    Edge,
    OnNodesChange,
    OnEdgesChange,
    OnConnect,
    NodeTypes,
    Handle,
    Position,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';
import { Bot, Play, Save, Plus, Settings } from 'lucide-react';
import { Button } from '@/components/ui/button';
import { cn } from '@/lib/utils';
import { API_URL } from '@/config/api';

// Custom Node Component
const AgentNode = ({ data }: { data: { label: string; role?: string } }) => {
    return (
        <div className="px-4 py-2 shadow-md rounded-md bg-card border border-border min-w-[150px]">
            <div className="flex items-center">
                <div className="rounded-full w-8 h-8 flex items-center justify-center bg-purple-500/20 text-purple-500 mr-2">
                    <Bot size={16} />
                </div>
                <div className="ml-2">
                    <div className="text-sm font-bold">{data.label}</div>
                    <div className="text-xs text-gray-500">{data.role || 'General Agent'}</div>
                </div>
            </div>
            <Handle type="target" position={Position.Top} className="w-16 !bg-muted-foreground" />
            <Handle type="source" position={Position.Bottom} className="w-16 !bg-muted-foreground" />
        </div>
    );
};

const nodeTypes: NodeTypes = {
    agent: AgentNode,
};

const initialNodes: Node[] = [
    {
        id: '1',
        type: 'agent',
        data: { label: 'Orchestrator', role: 'Supervisor' },
        position: { x: 250, y: 5 },
    },
];

const initialEdges: Edge[] = [];

export default function AgentFlowWidget() {
    const [nodes, setNodes] = useState<Node[]>(initialNodes);
    const [edges, setEdges] = useState<Edge[]>(initialEdges);
    const [isRunning, setIsRunning] = useState(false);

    const onNodesChange: OnNodesChange = useCallback(
        (changes) => setNodes((nds) => applyNodeChanges(changes, nds)),
        [],
    );
    const onEdgesChange: OnEdgesChange = useCallback(
        (changes) => setEdges((eds) => applyEdgeChanges(changes, eds)),
        [],
    );
    const onConnect: OnConnect = useCallback(
        (params) => setEdges((eds) => addEdge(params, eds)),
        [],
    );

    const addAgent = () => {
        const id = `agent-${nodes.length + 1}`;
        const newNode: Node = {
            id,
            type: 'agent',
            data: { label: `Agent ${nodes.length + 1}`, role: 'Worker' },
            position: { x: Math.random() * 400, y: Math.random() * 400 },
        };
        setNodes((nds) => [...nds, newNode]);
    };

    const runFlow = async () => {
        setIsRunning(true);
        try {
            // Prepare payload for backend
            const workflowPayload = {
                workflow: {
                    nodes: nodes.map(n => ({ id: n.id, type: n.type, data: n.data })),
                    edges: edges.map(e => ({ source: e.source, target: e.target }))
                }
            };

            // Call the backend handler (this matches the structure found in toolHandlers.ts)
            // Note: In a real app, this would be a proper API call or WebSocket event
            console.log('Running workflow:', workflowPayload);
            
            // Simulation
            await new Promise(resolve => setTimeout(resolve, 2000));
            
        } catch (error) {
            console.error('Flow execution failed:', error);
        } finally {
            setIsRunning(false);
        }
    };

    return (
        <div className="h-full flex flex-col bg-card/80 backdrop-blur-sm border border-border/50 rounded-lg overflow-hidden shadow-lg">
            {/* Header */}
            <div className="p-4 border-b border-border/50 bg-gradient-to-r from-purple-500/10 to-blue-500/5 flex items-center justify-between">
                <div className="flex items-center gap-3">
                    <div className="p-2 bg-gradient-to-br from-purple-500 to-blue-600 rounded-lg shadow-md">
                        <Bot className="w-5 h-5 text-white" />
                    </div>
                    <div>
                        <h2 className="font-bold text-foreground tracking-tight">Agent Flow Builder</h2>
                        <div className="flex items-center gap-2 text-[10px] font-mono text-muted-foreground">
                            <span className={cn("w-1.5 h-1.5 rounded-full", isRunning ? "bg-green-400 animate-pulse" : "bg-purple-400")} />
                            {isRunning ? "EXECUTING" : "DESIGN MODE"}
                        </div>
                    </div>
                </div>
                <div className="flex gap-2">
                     <Button size="sm" variant="outline" onClick={addAgent} className="h-8 gap-2">
                        <Plus size={14} /> Add Agent
                    </Button>
                    <Button size="sm" onClick={runFlow} disabled={isRunning} className="h-8 gap-2 bg-purple-600 hover:bg-purple-700">
                        <Play size={14} /> Run Flow
                    </Button>
                </div>
            </div>

            {/* Canvas */}
            <div className="flex-1 h-[500px] w-full bg-background/50">
                <ReactFlow
                    nodes={nodes}
                    edges={edges}
                    onNodesChange={onNodesChange}
                    onEdgesChange={onEdgesChange}
                    onConnect={onConnect}
                    nodeTypes={nodeTypes}
                    fitView
                >
                    <Background color="#888" gap={16} size={1} />
                    <Controls />
                </ReactFlow>
            </div>
        </div>
    );
}