components/AIBuilderModal.tsx

import React, { useState, useEffect, useRef } from 'react';
import { Sparkles, X, MessageSquare, Play, RefreshCw, AlertCircle, HardHat, Search, Wand2, CheckCircle2 } from 'lucide-react';
import { generateGraphWithAgents, AgentStatus } from '../services/geminiService';
import { Node } from 'reactflow';
import { NodeData } from '../types';

interface AIBuilderModalProps {
  isOpen: boolean;
  onClose: () => void;
  onApply: (nodes: any[], edges: any[], logMsg?: string) => void;
  currentNodes: Node<NodeData>[];
}

const AIBuilderModal: React.FC<AIBuilderModalProps> = ({ isOpen, onClose, onApply, currentNodes }) => {
  const [prompt, setPrompt] = useState('');
  const [isBuilding, setIsBuilding] = useState(false);
  const [error, setError] = useState<string | null>(null);
  
  const [agentStatus, setAgentStatus] = useState<AgentStatus>('idle');
  const [agentMessage, setAgentMessage] = useState('');

  // Reset state when opened
  useEffect(() => {
    if (isOpen) {
      setAgentStatus('idle');
      setAgentMessage('');
      setError(null);
    }
  }, [isOpen]);

  if (!isOpen) return null;

  const handleBuild = async () => {
    if (!prompt.trim()) return;
    setIsBuilding(true);
    setError(null);

    try {
      const result = await generateGraphWithAgents(prompt, currentNodes, (status, msg) => {
          setAgentStatus(status);
          setAgentMessage(msg);
      });
      
      if (result && result.nodes && result.edges) {
        // Ensure nodes have correct style defaults if missing
        const processedNodes = result.nodes.map((n: any) => {
            const data = n.data || {};
            const type = data.type || n.type || 'Identity'; // Fallback
            return {
              ...n,
              type: 'custom',
              data: {
                ...data,
                type: type,
                label: data.label || n.label || type, // Safe access
                params: data.params || {}
              }
            };
        });
        
        // Ensure edges have styling
        const processedEdges = result.edges.map((e: any) => ({
          ...e,
          animated: true, 
          style: { stroke: '#94a3b8' }
        }));

        // Small delay to let user see "Complete" state
        setTimeout(() => {
             onApply(processedNodes, processedEdges, "AI Architect generated new graph.");
             onClose();
             setIsBuilding(false);
        }, 1000);
      } else {
        throw new Error("Invalid response format from AI");
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : "Failed to generate architecture");
      setAgentStatus('error');
      setIsBuilding(false);
    }
  };

  const getStepStatus = (step: AgentStatus, current: AgentStatus) => {
    const order = ['idle', 'architect', 'critic', 'refiner', 'complete'];
    const stepIdx = order.indexOf(step);
    const currentIdx = order.indexOf(current);
    
    if (current === 'error') return 'text-slate-500';
    if (currentIdx > stepIdx) return 'text-emerald-400';
    if (currentIdx === stepIdx) return 'text-blue-400 animate-pulse';
    return 'text-slate-600';
  };

  const renderAgentStep = (step: AgentStatus, icon: React.ReactNode, label: string) => {
      const statusColor = getStepStatus(step, agentStatus);
      const isCurrent = agentStatus === step;
      const isDone = ['architect', 'critic', 'refiner', 'complete'].indexOf(agentStatus) > ['architect', 'critic', 'refiner'].indexOf(step);

      return (
        <div className={`flex items-center gap-3 p-3 rounded-lg border border-slate-800 ${isCurrent ? 'bg-slate-800/50 border-blue-500/30' : 'bg-slate-900'}`}>
            <div className={`${statusColor} transition-colors duration-300`}>
                {isDone ? <CheckCircle2 size={24} /> : icon}
            </div>
            <div className="flex-1">
                <div className={`font-semibold text-sm ${isCurrent ? 'text-blue-200' : 'text-slate-400'}`}>{label}</div>
                {isCurrent && (
                     <div className="text-xs text-slate-500 mt-1">{agentMessage}</div>
                )}
            </div>
            {isCurrent && <div className="w-2 h-2 rounded-full bg-blue-400 animate-ping" />}
        </div>
      );
  }

  return (
    <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm p-4">
      <div className="bg-slate-950 w-full max-w-lg rounded-xl border border-slate-700 shadow-2xl flex flex-col overflow-hidden animate-in fade-in zoom-in duration-200">
        
        <div className="flex items-center justify-between px-6 py-4 border-b border-slate-800 bg-gradient-to-r from-slate-900 to-slate-800">
          <h2 className="text-lg font-bold text-white flex items-center gap-2">
            <Sparkles className="text-purple-400" size={20} />
            AI Architect Team
          </h2>
          <button onClick={onClose} disabled={isBuilding} className="text-slate-400 hover:text-white transition-colors disabled:opacity-50">
            <X size={20} />
          </button>
        </div>

        <div className="p-6 space-y-4">
          {!isBuilding && agentStatus !== 'complete' ? (
            <>
                <div className="bg-purple-500/10 border border-purple-500/20 rounded-lg p-3 text-sm text-purple-200 flex gap-3">
                    <MessageSquare size={18} className="shrink-0 mt-0.5" />
                    <p>
                    Describe your model. Our AI Agents (Architect, Critic, and Refiner) will collaborate to build the best architecture for you.
                    </p>
                </div>

                <textarea
                    className="w-full h-32 bg-slate-900 border border-slate-700 rounded-lg p-3 text-slate-200 focus:ring-2 focus:ring-purple-500 outline-none resize-none placeholder-slate-600"
                    placeholder="e.g. Build a robust CNN for medical image segmentation with attention gates..."
                    value={prompt}
                    onChange={(e) => setPrompt(e.target.value)}
                />

                {error && (
                    <div className="flex items-center gap-2 text-red-400 text-sm bg-red-500/10 p-2 rounded border border-red-500/20">
                    <AlertCircle size={14} />
                    {error}
                    </div>
                )}

                <div className="flex justify-end pt-2">
                    <button
                    onClick={handleBuild}
                    disabled={!prompt.trim()}
                    className={`
                        flex items-center gap-2 px-6 py-2 rounded-lg font-medium text-white transition-all
                        ${!prompt.trim() 
                        ? 'bg-slate-800 cursor-not-allowed opacity-50' 
                        : 'bg-gradient-to-r from-purple-600 to-blue-600 hover:from-purple-500 hover:to-blue-500 shadow-lg shadow-purple-500/20'}
                    `}
                    >
                    <Play size={18} fill="currentColor" />
                    Start Agents
                    </button>
                </div>
            </>
          ) : (
            <div className="space-y-3 py-2">
                {renderAgentStep('architect', <HardHat size={24} />, "Architect Agent: Drafting Layout")}
                {renderAgentStep('critic', <Search size={24} />, "Critic Agent: Reviewing Design")}
                {renderAgentStep('refiner', <Wand2 size={24} />, "Refiner Agent: Finalizing Architecture")}
                
                {agentStatus === 'complete' && (
                     <div className="text-center text-emerald-400 font-bold mt-4 animate-pulse">
                        Architecture Generation Complete!
                     </div>
                )}
            </div>
          )}
        </div>
      </div>
    </div>
  );
};

export default AIBuilderModal;