src/components/AIArchitectPanel.tsx

'use client';

import React, { useState } from 'react';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Button } from '@/components/ui/button';
import { Textarea } from '@/components/ui/textarea';
import { Badge } from '@/components/ui/badge';
import { Alert, AlertDescription } from '@/components/ui/alert';
import { 
  Brain, 
  Zap, 
  Code2, 
  PlayCircle, 
  CheckCircle, 
  AlertCircle,
  Clock,
  Cpu,
  TrendingUp,
  Download,
  Database,
  FileText,
  Settings,
  Circle
} from 'lucide-react';
import quantumWeaverAPI from '@/lib/api-client';
import { Dialog, DialogContent, DialogHeader, DialogTitle } from '@/components/ui/dialog';
import { Separator } from '@/components/ui/separator';
import { useRouter } from 'next/navigation';
import UnifiedQuestionnaire from './UnifiedQuestionnaire';

interface AnalysisResult {
  task_type: string;
  domain: string;
  complexity: string;
  estimated_time: string;
  model_suggestions: string[];
  data_requirements: string[];
}

interface DataPlan {
  internet_datasets: Array<{
    name: string;
    description: string;
    url: string;
    download_size: string;
  }>;
  user_data_needed: boolean;
  user_data_instructions: string;
  pretrained_models: Array<{
    name: string;
    description: string;
  }>;
  preprocessing_steps: string[];
}

interface WorkflowResult {
  status: string;
  workflow_id: string;
  analysis: AnalysisResult;
  data_plan: DataPlan;
  training_pipeline: any;
  message: string;
  // Optional fields returned by backend we rely on
  jupyter_notebook?: {
    content: string;
    file_name?: string;
    visualizations?: string[];
    sections?: string[];
    interactive?: string[];
  } | null;
  solution?: any;
}

export default function AIArchitectPanel() {
  const router = useRouter();
  const [userRequest, setUserRequest] = useState('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [workflowResult, setWorkflowResult] = useState<WorkflowResult | null>(null);
  const [processingStep, setProcessingStep] = useState<string>('');
  const [isNotebookOpen, setIsNotebookOpen] = useState(false);
  const [trainingInfo, setTrainingInfo] = useState<{ status: 'idle'|'starting'|'started'|'error'; message?: string }>()
  const [workflowId, setWorkflowId] = useState<string | null>(null);
  const [workflowStatus, setWorkflowStatus] = useState<any | null>(null);
  const [polling, setPolling] = useState<any | null>(null);
  const [showQuestionnaire, setShowQuestionnaire] = useState(false);
  const [isSubmittingQuestionnaire, setIsSubmittingQuestionnaire] = useState(false);
  const [hasSubmittedQuestionnaire, setHasSubmittedQuestionnaire] = useState(false);
  const [questionnaireData, setQuestionnaireData] = useState<any>(null);
  const [questionnaireShownForWorkflow, setQuestionnaireShownForWorkflow] = useState<string | null>(null);

  // Smart API base detection with fallbacks
  const getApiBase = () => {
    if (typeof window === 'undefined') return process.env.NEXT_PUBLIC_TRAINING_API_BASE || 'http://localhost:9006';
    const envBase = process.env.NEXT_PUBLIC_TRAINING_API_BASE;
    const fallbacks = [
      envBase,
      'http://localhost:9006',
      'http://127.0.0.1:9006',
      'http://0.0.0.0:9006',
      (window.location && window.location.origin ? window.location.origin.replace(':9002', ':9006') : '')
    ].filter(Boolean);
    return fallbacks[0] as string;
  };
  const API_BASE = getApiBase();

  // Robust fetch that tries alternative backends with per-attempt timeout
  const smartFetch = async (path: string, init?: RequestInit, perAttemptTimeoutMs: number = 60000): Promise<Response> => {
    const candidates = [
      API_BASE,
      'http://127.0.0.1:9006',
      'http://0.0.0.0:9006',
      (typeof window !== 'undefined' && window.location ? window.location.origin.replace(':9002', ':9006') : '')
    ].filter(Boolean) as string[];

    let lastErr: any = null;
    for (const base of candidates) {
      const controller = new AbortController();
      const timer = setTimeout(() => controller.abort(), perAttemptTimeoutMs);
      try {
        const res = await fetch(`${base}${path}`, { ...(init || {}), signal: controller.signal });
        clearTimeout(timer);
        return res;
      } catch (e: any) {
        clearTimeout(timer);
        lastErr = e;
        continue;
      }
    }
    throw lastErr || new Error('Failed to fetch from all API bases');
  };

  const handleViewNotebook = () => {
    if (!workflowResult?.jupyter_notebook) {
      setError('Notebook not available from backend');
      return;
    }
    setIsNotebookOpen(true);
  };

  const handleStartTraining = async () => {
    if (!workflowResult) return;
    try {
      setTrainingInfo({ status: 'starting', message: 'Starting training...' });
      const res = await fetch(`${API_BASE}/api/ai-architect/start-training/${workflowResult.workflow_id}`, { method: 'POST' });
      const data = await res.json();
      if (!res.ok) {
        throw new Error(data?.message || 'Failed to start training');
      }
      setTrainingInfo({ status: 'started', message: data?.message || 'Training started' });
    } catch (e: any) {
      setTrainingInfo({ status: 'error', message: e?.message || 'Training failed' });
    }
  };

  const generateArchitecture = async () => {
    // Allow trigger even if input is empty; backend can prompt for details
    setIsGenerating(true);
    setError(null);
    setWorkflowResult(null);
    setProcessingStep('Analyzing your request...');
    setHasSubmittedQuestionnaire(false);
    setShowQuestionnaire(false);
    setQuestionnaireData(null);
    setQuestionnaireShownForWorkflow(null);
    
    try {
      console.log('🚀 Calling Ultimate AI Workflow Orchestrator via API client...');
      
      // Update progress steps
      setTimeout(() => setProcessingStep('Consulting MLE-STAR agents...'), 1000);
      setTimeout(() => setProcessingStep('Searching for optimal datasets...'), 2000);
      setTimeout(() => setProcessingStep('Generating training pipeline...'), 3000);
      
      // Hit the real backend to create and run the full workflow (A→Z)
      const res = await smartFetch('/api/ai-architect/create-workflow', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ user_request: userRequest, workflow_type: 'full_automation', priority: 'normal' })
      }, 60000);
      const data = await res.json();
      if (!res.ok) throw new Error(data?.detail || data?.message || 'Workflow creation failed');

      // Store workflowId and start polling status so the user sees the agents/stages progress
      setWorkflowId(data.workflow_id);
      // Seed immediate status so the progress card renders right away
      setWorkflowStatus({
        workflow_id: data.workflow_id,
        status: 'running',
        current_stage: 0,
        stages: [],
        progress_percentage: 0,
        created_at: new Date().toISOString()
      } as any);
      setProcessingStep('Workflow started. Tracking progress...');
             if (polling) clearInterval(polling);
      const tick = async () => {
        if (!data.workflow_id) return;
        try {
          const sRes = await fetch(`${API_BASE}/api/ai-architect/workflow-status/${data.workflow_id}?ts=${Date.now()}`, {
            headers: { 'Cache-Control': 'no-cache' },
            cache: 'no-store'
          });
          if (sRes.ok) {
            const status = await sRes.json();
            setWorkflowStatus(status);
            // Show questionnaire when available and not already submitted
            const dp = status?.solution?.data_plan;
            const hasQuestionnaire = dp && ((dp.user_upload_needed === true) || (Array.isArray(dp.questionnaires) && dp.questionnaires.length > 0));
            if (hasQuestionnaire && 
                !showQuestionnaire && 
                !hasSubmittedQuestionnaire && 
                status.status === 'waiting_for_user_data' &&
                questionnaireShownForWorkflow !== workflowId) {
              
              console.log('🎯 Fetching intelligent questionnaire for workflow:', workflowId);
              setQuestionnaireShownForWorkflow(workflowId);
              
              // Fetch questionnaire data
              fetchIntelligentQuestionnaire(workflowId).then(() => {
                setShowQuestionnaire(true);
              });
            }
            // When completed, stop polling and surface final artifacts
                          if (status.status === 'completed' || status.status === 'failed') {
                if (polling) clearInterval(polling);
                setProcessingStep(status.status === 'completed' ? 'Finalizing workflow...' : 'Workflow failed');
              setWorkflowResult({
                status: status.status === 'completed' ? 'success' : 'error',
                workflow_id: status.workflow_id,
                analysis: {
                  task_type: status?.solution?.training_pipeline?.architecture?.task_type || 'unspecified',
                  domain: status?.solution?.training_pipeline?.architecture?.domain || 'unspecified',
                  complexity: 'unspecified',
                  estimated_time: 'unspecified',
                  model_suggestions: [] as string[],
                  data_requirements: [] as string[]
                },
                data_plan: status?.solution?.data_plan || { internet_datasets: [], user_data_needed: false, user_data_instructions: '', pretrained_models: [], preprocessing_steps: [] },
                training_pipeline: status?.solution?.training_pipeline || {},
                message: 'Workflow completed',
                jupyter_notebook: status?.solution?.jupyter_notebook || null,
                solution: status?.solution || {}
              });
              setProcessingStep('');
            }
          }
        } catch (e) {
          // swallow polling errors
        }
      };
      const intervalId = setInterval(tick, 1000);
      setPolling(intervalId);
      // fire immediately so user sees first stage
      tick();

      
      // Keep input so user can tweak and resubmit; do not clear request text
      // setUserRequest('');
     
    } catch (error: any) {
      console.error('❌ AI Architect connection failed:', error);
      setError(`Failed to connect to AI Workflow Orchestrator: ${error.message}`);
      setProcessingStep('');
    } finally {
      // Ensure the button re-enables even on failures/timeouts
      setIsGenerating(false);
      // Focus button again for immediate retry UX
      try { (document.querySelector('#genAgentsBtn') as HTMLButtonElement)?.focus(); } catch (err) {}
    }
  };

  const clearAll = () => {
    setUserRequest('');
    setError(null);
    setWorkflowResult(null);
    setWorkflowStatus(null);
    setShowQuestionnaire(false);
    setHasSubmittedQuestionnaire(false);
    setQuestionnaireData(null);
    setQuestionnaireShownForWorkflow(null);
    if (polling) clearInterval(polling);
  };

  // Removed old submitDataModal - now using handleQuestionnaireSubmit

  // Removed old uploadFile function - file upload handled in UnifiedQuestionnaire

  const fetchIntelligentQuestionnaire = async (workflowId: string) => {
    try {
      console.log('📡 Fetching questionnaire data for workflow:', workflowId);
      const response = await smartFetch(`/api/ai-architect/questionnaire/${workflowId}`);
      if (response.ok) {
        const data = await response.json();
        console.log('📥 Received questionnaire response:', data);
        if (data.status === 'available' && data.questionnaire) {
          console.log('✅ Setting questionnaire data:', data.questionnaire);
          setQuestionnaireData(data.questionnaire);
          return true;
        } else {
          console.warn('⚠️ No questionnaire available, using defaults');
          setQuestionnaireData(null);
          return false;
        }
      } else {
        console.error('❌ Failed to fetch questionnaire, status:', response.status);
        setQuestionnaireData(null);
        return false;
      }
    } catch (error) {
      console.error('❌ Failed to fetch intelligent questionnaire:', error);
      setQuestionnaireData(null);
      return false;
    }
  };

  const handleQuestionnaireSubmit = async (answers: Record<string, any>) => {
    if (!workflowId) return;
    
    try {
      setIsSubmittingQuestionnaire(true);
      
      const response = await fetch(`${API_BASE}/api/ai-architect/submit-questionnaire`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          workflow_id: workflowId,
          answers: answers
        })
      });
      
      if (response.ok) {
        const result = await response.json();
        console.log('✅ Questionnaire submitted successfully:', result);
        
        // Close questionnaire modal and mark as submitted
        setShowQuestionnaire(false);
        setHasSubmittedQuestionnaire(true);
        setQuestionnaireData(null);
        
        console.log('🔄 Continuing workflow with answers...');
        // Continue the workflow
        const continueResponse = await fetch(`${API_BASE}/api/ai-architect/continue-workflow/${workflowId}`, {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify(answers)
        });
        
        if (continueResponse.ok) {
          const continueResult = await continueResponse.json();
          console.log('✅ Workflow continued successfully:', continueResult);
        } else {
          console.error('❌ Failed to continue workflow:', continueResponse.statusText);
        }
        
      } else {
        throw new Error(`Failed to submit questionnaire: ${response.statusText}`);
      }
    } catch (error) {
      console.error('❌ Failed to submit questionnaire:', error);
      alert('Failed to submit questionnaire. Please try again.');
    } finally {
      setIsSubmittingQuestionnaire(false);
    }
  };

  React.useEffect(() => {
    const checkStatus = async () => {
      if (workflowId) {
        try {
          const sRes = await fetch(`${API_BASE}/api/ai-architect/workflow-status/${workflowId}?ts=${Date.now()}`, {
            headers: { 'Cache-Control': 'no-cache' },
            cache: 'no-store'
          });
          if (sRes.ok) {
            const status = await sRes.json();
            setWorkflowStatus(status);
            
            // Check if we need to show the questionnaire (keep it open once shown until user acts)
            const dp = status?.solution?.data_plan;
            const hasQuestionnaire = dp && ((dp.user_upload_needed === true) || (Array.isArray(dp.questionnaires) && dp.questionnaires.length > 0));
            if (hasQuestionnaire && !hasSubmittedQuestionnaire) {
              setShowQuestionnaire(true);
            }
          }
        } catch (error) {
          console.error('Error checking workflow status:', error);
        }
      }
    };

    if (workflowId) {
      checkStatus();
      const interval = setInterval(checkStatus, 3000);
      return () => clearInterval(interval);
    }
  }, [workflowId, workflowResult]);

  return (
    <div className="w-full min-h-0 flex flex-col space-y-6 p-6 overflow-y-auto cyber-scrollbar">
      {/* AI Architect Header */}
      <Card className="border-blue-500/20 bg-gradient-to-r from-blue-900/20 to-purple-900/20">
        <CardHeader>
          <CardTitle className="flex items-center justify-between">
            <div className="flex items-center space-x-2">
              <Brain className="w-8 h-8 text-blue-400" />
              <span className="text-2xl">ZPE AI Architect V2</span>
              <Badge variant="secondary" className="bg-blue-500/20 text-blue-300">
                Full Automation Mode + Dataset Architect
              </Badge>
            </div>
            <div className="flex items-center space-x-2">
              <Badge className="bg-[#00ffe7]/20 text-[#00ffe7] border-[#00ffe7]/30">
                Ultimate AI Workflow ACTIVE
              </Badge>
              <Button variant="outline" size="sm" onClick={clearAll}>
                Clear All
              </Button>
            </div>
          </CardTitle>
        </CardHeader>
      </Card>

      {/* Request Input */}
      <Card>
        <CardHeader>
          <CardTitle className="flex items-center space-x-2">
            <Brain className="w-5 h-5" />
            <span>Natural Language Architecture Request</span>
          </CardTitle>
        </CardHeader>
        <CardContent className="space-y-4">
          <Textarea
            placeholder={`Describe the complete AI system you want to build (full automation from data to deployment)...

Examples:
• 'I need an AI to detect anomalies in manufacturing sensor data and deploy it to production'
• 'Create a complete voice recognition system for smart home control with monitoring'
• 'Build an end-to-end computer vision model to identify plant diseases with visualization'
• 'I want an AI that recognizes my face from my webcam and bypasses my password to log in'`}
            value={userRequest}
            onChange={(e) => setUserRequest(e.target.value)}
            rows={8}
            className="resize-none overflow-y-auto cyber-scrollbar"
          />
          
          <div className="flex items-center space-x-4">
            <Button 
              id="genAgentsBtn"
              type="button"
              onClick={(e) => { e.preventDefault(); e.stopPropagation(); generateArchitecture(); }}
              disabled={isGenerating}
              className="relative z-50 pointer-events-auto bg-gradient-to-r from-blue-600 to-purple-600 hover:from-blue-700 hover:to-purple-700"
            >
              {isGenerating ? (
                <>
                  <Zap className="w-4 h-4 mr-2 animate-spin" />
                  {processingStep || 'AI Agents Analyzing...'}
                </>
              ) : (
                <>
                  <Brain className="w-4 h-4 mr-2" />
                  Generate with AI Agents
                </>
              )}
            </Button>
            <div className="flex items-center space-x-2 text-sm text-muted-foreground">
              <Cpu className="w-4 h-4" />
              <span>ZPE Consciousness Enhanced</span>
            </div>
          </div>

          {/* Processing Status */}
          {isGenerating && processingStep && (
            <Alert className="border-blue-500/20 bg-blue-500/10">
              <Zap className="h-4 w-4 animate-pulse" />
              <AlertDescription className="text-blue-400">
                <div className="flex items-center space-x-2">
                  <span>{processingStep}</span>
                  <div className="flex space-x-1">
                    <div className="w-2 h-2 bg-blue-400 rounded-full animate-bounce" style={{animationDelay: '0ms'}}></div>
                    <div className="w-2 h-2 bg-blue-400 rounded-full animate-bounce" style={{animationDelay: '150ms'}}></div>
                    <div className="w-2 h-2 bg-blue-400 rounded-full animate-bounce" style={{animationDelay: '300ms'}}></div>
                  </div>
                </div>
              </AlertDescription>
            </Alert>
          )}

          {error && (
            <Alert className="border-red-500/20 bg-red-500/10">
              <AlertCircle className="h-4 w-4" />
              <AlertDescription className="text-red-400">
                {error}
              </AlertDescription>
            </Alert>
          )}
        </CardContent>
      </Card>

                {/* Live Agent/Stage Progress (visible immediately) */}
          {workflowStatus && (
            <Card className="border-cyan-500/20 bg-cyan-500/5">
              <CardHeader>
                <CardTitle className="flex items-center space-x-2 text-cyan-400">
                  <Clock className="w-5 h-5" />
                  <span>Agent Workflow Progress</span>
                </CardTitle>
              </CardHeader>
              <CardContent className="space-y-3">
                <div className="text-xs text-muted-foreground">Workflow ID: {workflowId}</div>
                <div className="text-sm">Status: <span className="font-semibold">{workflowStatus.status}</span></div>
                <div className="w-full h-2 bg-cyan-500/10 rounded">
                  <div className="h-2 bg-cyan-400 rounded transition-all duration-500" style={{ width: `${Math.min(100, workflowStatus.progress_percentage || 0)}%` }} />
                </div>
                <div className="space-y-3 max-h-64 overflow-auto cyber-scrollbar">
                  {(workflowStatus.stages || []).map((s: any, i: number) => (
                    <div key={i} className="space-y-1">
                      <div className="flex items-center justify-between text-sm">
                        <div className="flex items-center gap-2">
                          {s.status === 'completed' ? (
                            <CheckCircle className="w-4 h-4 text-green-400" />
                          ) : s.status === 'running' ? (
                            <div className="w-4 h-4 rounded-full border-2 border-cyan-400 border-t-transparent animate-spin" />
                          ) : (
                            <Circle className="w-4 h-4 text-muted-foreground" />
                          )}
                          <span className={s.status === 'running' ? 'text-cyan-400 font-medium' : ''}>{s.name}</span>
                        </div>
                        <span className="text-xs text-muted-foreground">{s.completed_at?.slice(11,19)}</span>
                      </div>
                      
                      {/* Show current substep if stage is running */}
                      {s.status === 'running' && s.current_substep && (
                        <div className="ml-6 pl-2 border-l border-cyan-500/30">
                          <div className="flex items-center gap-2 text-xs text-cyan-400/80">
                            <div className="w-1.5 h-1.5 bg-cyan-400 rounded-full animate-pulse" />
                            <span>{s.current_substep}</span>
                          </div>
                        </div>
                      )}
                      
                      {/* Show completed substeps */}
                      {s.substeps && s.substeps.length > 0 && (
                        <div className="ml-6 pl-2 border-l border-gray-500/30 space-y-1">
                          {s.substeps.map((sub: any, si: number) => (
                            <div key={si} className="flex items-center gap-2 text-xs text-muted-foreground">
                              <CheckCircle className="w-3 h-3 text-green-400/60" />
                              <span>{sub.name}</span>
                              {sub.found !== undefined && <span className="text-cyan-400">({sub.found} found)</span>}
                            </div>
                          ))}
                        </div>
                      )}
                    </div>
                  ))}
                </div>
              </CardContent>
            </Card>
          )}

          {/* Workflow Results */}
          {workflowResult && (
            <>

              {/* Analysis Results */}
          <Card className="border-green-500/20 bg-green-500/5">
            <CardHeader>
              <CardTitle className="flex items-center space-x-2 text-green-400">
                <CheckCircle className="w-5 h-5" />
                <span>AI Analysis Complete</span>
              </CardTitle>
            </CardHeader>
            <CardContent className="space-y-4">
              {/* Concise A–Z Plan Summary */}
              <div className="text-sm text-muted-foreground">
                End-to-end plan generated by MLE-STAR agents and your PyTorch architect: Data discovery → Data plan → Advanced prep & validation → Notebook generation → Training pipeline → Deployment hooks.
              </div>
              <div className="grid grid-cols-2 gap-4">
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Task Type</label>
                  <p className="text-lg font-semibold">{workflowResult.analysis?.task_type || 'Not specified'}</p>
                </div>
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Domain</label>
                  <p className="text-lg font-semibold">{workflowResult.analysis?.domain || 'Not specified'}</p>
                </div>
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Complexity</label>
                  <p className="text-lg font-semibold">{workflowResult.analysis?.complexity || 'Not specified'}</p>
                </div>
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Estimated Time</label>
                  <p className="text-lg font-semibold">{workflowResult.analysis?.estimated_time || 'Not specified'}</p>
                </div>
              </div>
              
              <div>
                <label className="text-sm font-medium text-muted-foreground">Recommended Models</label>
                <div className="flex flex-wrap gap-2 mt-2">
                  {(workflowResult.analysis?.model_suggestions || []).map((model, idx) => (
                    <Badge key={idx} className="bg-blue-500/20 text-blue-300">
                      {model}
                    </Badge>
                  ))}
                </div>
              </div>
            </CardContent>
          </Card>

          {/* Data Plan */}
          <Card className="border-purple-500/20 bg-purple-500/5">
            <CardHeader>
              <CardTitle className="flex items-center space-x-2 text-purple-400">
                <Database className="w-5 h-5" />
                <span>Intelligent Data Plan</span>
              </CardTitle>
            </CardHeader>
            <CardContent className="space-y-4">
              {/* Internet Datasets */}
              {(workflowResult.data_plan?.internet_datasets?.length || 0) > 0 && (
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Available Datasets</label>
                  <div className="space-y-2 mt-2">
                    {(workflowResult.data_plan?.internet_datasets || []).map((dataset, idx) => (
                      <div key={idx} className="p-3 border border-purple-500/20 rounded-lg bg-purple-500/5">
                        <div className="flex justify-between items-start">
                          <div>
                            <h4 className="font-semibold">{dataset.name}</h4>
                            <p className="text-sm text-muted-foreground">{dataset.description}</p>
                            <p className="text-xs text-purple-400">Size: {dataset.download_size}</p>
                          </div>
                          <Button variant="outline" size="sm">
                            <Download className="w-4 h-4" />
                          </Button>
                        </div>
                      </div>
                    ))}
                  </div>
                </div>
              )}

              {/* User Data Requirements */}
              {workflowResult.data_plan?.user_data_needed && (
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Personal Data Required</label>
                  <Alert className="border-yellow-500/20 bg-yellow-500/10 mt-2">
                    <AlertCircle className="h-4 w-4" />
                    <AlertDescription className="text-yellow-400 whitespace-pre-line">
                      {workflowResult.data_plan?.user_data_instructions || 'No specific instructions'}
                    </AlertDescription>
                  </Alert>
                </div>
              )}

              {/* Pretrained Models */}
              {(workflowResult.data_plan?.pretrained_models?.length || 0) > 0 && (
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Pretrained Models</label>
                  <div className="flex flex-wrap gap-2 mt-2">
                    {(workflowResult.data_plan?.pretrained_models || []).map((model, idx) => (
                      <Badge key={idx} className="bg-green-500/20 text-green-300">
                        {model.name}
                      </Badge>
                    ))}
                  </div>
                </div>
              )}

              {/* Preprocessing Steps */}
              <div>
                <label className="text-sm font-medium text-muted-foreground">Preprocessing Pipeline</label>
                <div className="space-y-1 mt-2">
                  {(workflowResult.data_plan?.preprocessing_steps || []).map((step, idx) => (
                    <div key={idx} className="flex items-center space-x-2 text-sm">
                      <Settings className="w-4 h-4 text-blue-400" />
                      <span>{step}</span>
                    </div>
                  ))}
                </div>
              </div>
            </CardContent>
          </Card>

          {/* Training Pipeline */}
          <Card className="border-orange-500/20 bg-orange-500/5">
            <CardHeader>
              <CardTitle className="flex items-center space-x-2 text-orange-400">
                <Code2 className="w-5 h-5" />
                <span>Training Pipeline Generated</span>
              </CardTitle>
            </CardHeader>
            <CardContent className="space-y-4">
              <div className="grid grid-cols-2 gap-4">
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Model Architecture</label>
                  <p className="text-lg font-semibold">
                    {typeof workflowResult.training_pipeline?.architecture === 'object' 
                      ? workflowResult.training_pipeline?.architecture?.model_name || 'Custom Architecture'
                      : workflowResult.training_pipeline?.architecture || 'Not specified'}
                  </p>
                </div>
                <div>
                  <label className="text-sm font-medium text-muted-foreground">Backbone</label>
                  <p className="text-lg font-semibold">{workflowResult.training_pipeline?.backbone || 'Not specified'}</p>
                </div>
              </div>
              
              {/* Training Pipeline Generated section buttons */}
              <div className="flex items-center gap-3">
                <Button
                  variant="secondary"
                  onClick={() => {
                    if (workflowId) {
                      window.open(`${API_BASE}/api/ai-architect/notebook-html/${workflowId}`, '_blank');
                    }
                  }}
                >
                  View Jupyter Notebook
                </Button>
                <Button
                  className="bg-green-700 hover:bg-green-800"
                  onClick={async () => {
                    if (!workflowId) return;
                    const res = await fetch(`${API_BASE}/api/ai-architect/start-training/${workflowId}`, { method: 'POST' });
                    if (!res.ok) {
                      const t = await res.text();
                      setError(`Failed to start training: ${t}`);
                    } else {
                      setProcessingStep('Starting training...');
                    }
                  }}
                >
                  Start Training
                </Button>
                <Button
                  variant="outline"
                  onClick={() => { if (workflowId) window.location.href = `/deep-learning/train?workflow=${workflowId}`; }}
                >
                  Go to Train Page
                </Button>
                <Button
                  variant="secondary"
                  onClick={() => { if (workflowId) window.open(`${API_BASE}/api/ai-architect/artifacts/${workflowId}/zip`, '_blank'); }}
                >
                  Download Artifacts
                </Button>
                {/* New: App Builder & Deployer */}
                <Button
                  variant="secondary"
                  onClick={async () => {
                    try {
                      if (!workflowId) return;
                      setProcessingStep('Preparing app artifacts...');
                      const res = await fetch(`${API_BASE}/api/app-builder/prepare-artifacts/${workflowId}`, { method: 'POST' });
                      const data = await res.json();
                      if (!res.ok) throw new Error(data?.error || 'Prepare artifacts failed');
                      setProcessingStep(`App prepared at ${data.app_dir}`);
                    } catch (e: any) {
                      setError(e.message);
                      setProcessingStep('');
                    }
                  }}
                >
                  Prepare App
                </Button>
                <Button
                  variant="outline"
                  onClick={async () => {
                    try {
                      if (!workflowId) return;
                      setProcessingStep('Building container image...');
                      const res = await fetch(`${API_BASE}/api/deployer/build-image/${workflowId}`, { method: 'POST' });
                      const data = await res.json();
                      if (!res.ok) throw new Error(data?.error || 'Build image failed');
                      setProcessingStep(`Image build status: ${data.status}${data.image ? ' (' + data.image + ')' : ''}`);
                    } catch (e: any) {
                      setError(e.message);
                      setProcessingStep('');
                    }
                  }}
                >
                  Build Image
                </Button>
                <Button
                  variant="default"
                  onClick={async () => {
                    try {
                      if (!workflowId) return;
                      setProcessingStep('Starting local container...');
                      const res = await fetch(`${API_BASE}/api/deployer/run-local/${workflowId}`, { method: 'POST' });
                      const data = await res.json();
                      if (!res.ok) throw new Error(data?.error || 'Run local failed');
                      setProcessingStep(`Local server: http://localhost:${(data.ports||[8080])[0]}/health`);
                    } catch (e: any) {
                      setError(e.message);
                      setProcessingStep('');
                    }
                  }}
                >
                  Run Locally
                </Button>
              </div>
              {trainingInfo?.status && trainingInfo.status !== 'idle' && (
                <div className="mt-3 text-sm">
                  <span className={
                    trainingInfo.status === 'started' ? 'text-green-400' : trainingInfo.status === 'error' ? 'text-red-400' : 'text-muted-foreground'
                  }>
                    {trainingInfo.message}
                  </span>
                </div>
              )}
              
              <Alert className="border-orange-500/20 bg-orange-500/10">
                <CheckCircle className="h-4 w-4" />
                <AlertDescription className="text-orange-400">
                  Complete training pipeline generated with {Object.keys(workflowResult.training_pipeline || {}).length} components: 
                  preprocessing, training, evaluation, and deployment code.
                </AlertDescription>
              </Alert>
            </CardContent>
          </Card>
        </>
      )}

      {/* Dataset Architect Integration Notice */}
      <Card className="border-[#00ffe7]/30 bg-[#00ffe7]/5">
        <CardHeader>
          <CardTitle className="flex items-center space-x-2 text-[#00ffe7]">
            <CheckCircle className="w-5 h-5" />
            <span>Ultimate AI Workflow Integration</span>
          </CardTitle>
        </CardHeader>
        <CardContent>
          <div className="space-y-3">
            <p className="text-[#00ffe7]/80">
              The Ultimate AI Workflow Orchestrator is now FULLY ACTIVE! It provides complete automation from 
              prompt analysis to deployment-ready models with real datasets and Jupyter notebooks.
            </p>
            <div className="flex items-center space-x-2">
              <Badge className="bg-[#39ff14]/20 text-[#39ff14] border-[#39ff14]/30">
                <CheckCircle className="w-3 h-3 mr-1" />
                Smart Prompt Analysis
              </Badge>
              <Badge className="bg-[#39ff14]/20 text-[#39ff14] border-[#39ff14]/30">
                <CheckCircle className="w-3 h-3 mr-1" />
                Real Dataset Discovery
              </Badge>
              <Badge className="bg-[#39ff14]/20 text-[#39ff14] border-[#39ff14]/30">
                <CheckCircle className="w-3 h-3 mr-1" />
                Complete Training Pipeline
              </Badge>
            </div>
          </div>
        </CardContent>
      </Card>

      {/* Notebook Modal */}
      <Dialog open={isNotebookOpen} onOpenChange={setIsNotebookOpen}>
        <DialogContent className="max-w-[90vw] w-[90vw] max-h-[90vh] h-[90vh] overflow-hidden">
          <DialogHeader>
            <DialogTitle className="flex items-center gap-2">
              <FileText className="w-4 h-4" />
              {workflowResult?.jupyter_notebook?.file_name || 'Generated Notebook'}
            </DialogTitle>
          </DialogHeader>
          {workflowId ? (
            <iframe
              title="Notebook Viewer"
              src={`${API_BASE}/api/ai-architect/notebook-html/${workflowId}`}
              className="w-full h-[80vh] rounded border"
            />
          ) : (
            <div className="text-sm text-muted-foreground">Notebook content not available</div>
          )}
        </DialogContent>
      </Dialog>

      {/* Unified Questionnaire Modal */}
      <UnifiedQuestionnaire
        isOpen={showQuestionnaire}
        onClose={() => {
          console.log('🚪 Closing questionnaire modal');
          setShowQuestionnaire(false);
          setHasSubmittedQuestionnaire(true);
        }}
        onSubmit={handleQuestionnaireSubmit}
        isSubmitting={isSubmittingQuestionnaire}
        workflowId={workflowId}
        questionnaireData={questionnaireData}
      />


    </div>
  );
}