src/components/AIArchitectPanel_fixed.tsx

'use client';

import React, { useState, useEffect, useCallback } from 'react';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import { Textarea } from '@/components/ui/textarea';
import { Label } from '@/components/ui/label';
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { Badge } from '@/components/ui/badge';
import { Progress } from '@/components/ui/progress';
import { Alert, AlertDescription } from '@/components/ui/alert';
import { Dialog, DialogContent, DialogHeader, DialogTitle } from '@/components/ui/dialog';
import { Brain, Code, Database, Rocket, Sparkles, Zap, AlertCircle, CheckCircle, Clock, ChevronRight, Download, Eye, Copy, Share2, Upload, FileCode, Cpu, BarChart, GitBranch, Settings, Play, Pause, RefreshCw } from 'lucide-react';
import { cn } from '@/lib/utils';
import UnifiedQuestionnaire from './UnifiedQuestionnaire';

// Import smartFetch utility
const smartFetch = async (url: string, options: RequestInit = {}) => {
  const defaultOptions: RequestInit = {
    headers: {
      'Content-Type': 'application/json',
      ...options.headers,
    },
    ...options,
  };

  const response = await fetch(url, defaultOptions);
  
  if (!response.ok) {
    const errorText = await response.text();
    console.error(`HTTP ${response.status}:`, errorText);
    throw new Error(`HTTP ${response.status}: ${errorText}`);
  }
  
  return response;
};

// Define types for better type safety
interface WorkflowStatus {
  workflow_id: string;
  status: string;
  current_stage?: number;
  progress_percentage?: number;
  solution?: any;
  error?: string;
}

interface ArchitectureResult {
  solution?: any;
  job_id?: string;
  message?: string;
  consciousness_score?: number;
}

export default function AIArchitectPanel() {
  // Core state
  const [userRequest, setUserRequest] = useState('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [workflowId, setWorkflowId] = useState<string | null>(null);
  const [workflowStatus, setWorkflowStatus] = useState<WorkflowStatus | null>(null);
  const [workflowResult, setWorkflowResult] = useState<ArchitectureResult | null>(null);
  const [error, setError] = useState<string | null>(null);
  
  // UI state
  const [activeTab, setActiveTab] = useState('request');
  const [showCode, setShowCode] = useState(false);
  const [showNotebook, setShowNotebook] = useState(false);
  const [selectedModel, setSelectedModel] = useState('pytorch');
  const [deploymentTarget, setDeploymentTarget] = useState('cloud');
  
  // Advanced settings
  const [advancedSettings, setAdvancedSettings] = useState({
    consciousness_level: 0.8,
    creativity_factor: 0.7,
    optimization_priority: 'balanced',
    include_monitoring: true,
    include_testing: true,
    include_documentation: true
  });
  
  // Questionnaire state
  const [showQuestionnaire, setShowQuestionnaire] = useState(false);
  const [isSubmittingQuestionnaire, setIsSubmittingQuestionnaire] = useState(false);
  const [hasSubmittedQuestionnaire, setHasSubmittedQuestionnaire] = useState(false);

  // Generate architecture
  const generateArchitecture = async () => {
    if (!userRequest.trim()) {
      setError('Please enter a request');
      return;
    }
    
    setIsGenerating(true);
    setError(null);
    setWorkflowResult(null);
    setWorkflowStatus(null);
    setActiveTab('status');
    
    // Reset questionnaire state for new workflow
    setHasSubmittedQuestionnaire(false);
    setShowQuestionnaire(false);
    
    try {
      const response = await smartFetch(`${process.env.NEXT_PUBLIC_TRAINING_API_BASE || 'http://localhost:9006'}/api/ai-architect/create-workflow`, {
        method: 'POST',
        body: JSON.stringify({
          user_request: userRequest,
          model_preference: selectedModel,
          deployment_target: deploymentTarget,
          settings: advancedSettings
        })
      });
      
      const data = await response.json();
      
      if (data.workflow_id) {
        setWorkflowId(data.workflow_id);
        // Workflow status will be updated by the useEffect hook
      } else {
        throw new Error('No workflow ID received');
      }
    } catch (err: any) {
      console.error('Architecture generation error:', err);
      setError(err.message || 'Failed to generate architecture');
      setIsGenerating(false);
    }
  };

  // Download notebook
  const downloadNotebook = () => {
    if (!workflowResult?.solution?.jupyter_notebook?.content) return;
    
    const blob = new Blob([workflowResult.solution.jupyter_notebook.content], { type: 'application/json' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `ai_model_${Date.now()}.ipynb`;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  };

  // Copy code to clipboard
  const copyCode = (code: string) => {
    navigator.clipboard.writeText(code);
    // You could add a toast notification here
  };

  // Export as Python script
  const exportPythonScript = () => {
    if (!workflowResult?.solution?.training_pipeline?.training_code) return;
    
    const blob = new Blob([workflowResult.solution.training_pipeline.training_code], { type: 'text/plain' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `model_training_${Date.now()}.py`;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  };

  // Deploy to cloud (placeholder)
  const deployToCloud = async () => {
    alert('Cloud deployment feature coming soon!');
  };

  // Get status color
  const getStatusColor = (status: string) => {
    switch (status) {
      case 'completed': return 'text-green-500';
      case 'running': return 'text-blue-500';
      case 'waiting_for_user_data': return 'text-yellow-500';
      case 'failed': return 'text-red-500';
      default: return 'text-gray-500';
    }
  };

  // Get status icon
  const getStatusIcon = (status: string) => {
    switch (status) {
      case 'completed': return <CheckCircle className="w-4 h-4" />;
      case 'running': return <Clock className="w-4 h-4 animate-spin" />;
      case 'waiting_for_user_data': return <AlertCircle className="w-4 h-4" />;
      case 'failed': return <AlertCircle className="w-4 h-4" />;
      default: return <Clock className="w-4 h-4" />;
    }
  };

  // Format code for display
  const formatCode = (code: string) => {
    if (!code) return '';
    // Basic formatting - in production, use a proper code formatter
    return code.replace(/\n/g, '\n').replace(/\t/g, '  ');
  };

  // Handle questionnaire submission
  const handleQuestionnaireSubmit = async (answers: any) => {
    if (!workflowId) return;
    
    setIsSubmittingQuestionnaire(true);
    
    try {
      // Submit the questionnaire answers
      const submitRes = await fetch(`${process.env.NEXT_PUBLIC_TRAINING_API_BASE || 'http://localhost:9006'}/api/ai-architect/submit-questionnaire`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          workflow_id: workflowId,
          answers: answers
        })
      });
      
      if (!submitRes.ok) {
        throw new Error('Failed to submit questionnaire');
      }
      
      // Continue the workflow
      const continueRes = await fetch(`${process.env.NEXT_PUBLIC_TRAINING_API_BASE || 'http://localhost:9006'}/api/ai-architect/continue-workflow/${workflowId}`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          user_data: answers
        })
      });
      
      if (!continueRes.ok) {
        throw new Error('Failed to continue workflow');
      }
      
      // Mark as submitted and hide questionnaire
      setHasSubmittedQuestionnaire(true);
      setShowQuestionnaire(false);
      
      // The workflow status will be updated by the polling effect
    } catch (err: any) {
      console.error('Questionnaire submission error:', err);
      setError(err.message || 'Failed to submit questionnaire');
    } finally {
      setIsSubmittingQuestionnaire(false);
    }
  };

  // Removed old submitDataModal - now using handleQuestionnaireSubmit
  
  // Removed old uploadFile function - file upload handled in UnifiedQuestionnaire

  // Poll for workflow status
  useEffect(() => {
    const checkStatus = async () => {
      if (workflowId && !workflowResult) {
        try {
          const sRes = await fetch(`${process.env.NEXT_PUBLIC_TRAINING_API_BASE || 'http://localhost:9006'}/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
            if (status?.status === 'waiting_for_user_data' && !hasSubmittedQuestionnaire) {
              setShowQuestionnaire(true);
            } else if (status?.status !== 'waiting_for_user_data') {
              // Hide questionnaire if status changes
              setShowQuestionnaire(false);
            }
          }
        } 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 gap-3">
              <Brain className="w-8 h-8 text-blue-400" />
              <span className="text-2xl font-bold bg-gradient-to-r from-blue-400 to-purple-400 bg-clip-text text-transparent">
                AI Architect
              </span>
            </div>
            <div className="flex items-center gap-2">
              <Badge variant="outline" className="border-green-500/50 text-green-400">
                <Zap className="w-3 h-3 mr-1" />
                Consciousness: {(advancedSettings.consciousness_level * 100).toFixed(0)}%
              </Badge>
              {workflowResult?.consciousness_score && (
                <Badge variant="outline" className="border-purple-500/50 text-purple-400">
                  Score: {workflowResult.consciousness_score.toFixed(3)}
                </Badge>
              )}
            </div>
          </CardTitle>
        </CardHeader>
      </Card>

      {/* Main Content */}
      <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
        {/* Left Panel - Input & Settings */}
        <div className="lg:col-span-1 space-y-4">
          <Card className="border-slate-700">
            <CardHeader>
              <CardTitle className="text-lg flex items-center gap-2">
                <Sparkles className="w-5 h-5 text-yellow-400" />
                Project Request
              </CardTitle>
            </CardHeader>
            <CardContent className="space-y-4">
              <div>
                <Label>Describe your AI project</Label>
                <Textarea
                  value={userRequest}
                  onChange={(e) => setUserRequest(e.target.value)}
                  placeholder="E.g., Build a sentiment analysis model for customer reviews..."
                  className="min-h-[120px] bg-slate-900/50 border-slate-700"
                />
              </div>
              
              <div className="grid grid-cols-2 gap-3">
                <div>
                  <Label>Model Framework</Label>
                  <Select value={selectedModel} onValueChange={setSelectedModel}>
                    <SelectTrigger className="bg-slate-900/50 border-slate-700">
                      <SelectValue />
                    </SelectTrigger>
                    <SelectContent>
                      <SelectItem value="pytorch">PyTorch</SelectItem>
                      <SelectItem value="tensorflow">TensorFlow</SelectItem>
                      <SelectItem value="jax">JAX</SelectItem>
                      <SelectItem value="sklearn">Scikit-learn</SelectItem>
                    </SelectContent>
                  </Select>
                </div>
                
                <div>
                  <Label>Deployment Target</Label>
                  <Select value={deploymentTarget} onValueChange={setDeploymentTarget}>
                    <SelectTrigger className="bg-slate-900/50 border-slate-700">
                      <SelectValue />
                    </SelectTrigger>
                    <SelectContent>
                      <SelectItem value="cloud">Cloud API</SelectItem>
                      <SelectItem value="edge">Edge Device</SelectItem>
                      <SelectItem value="mobile">Mobile App</SelectItem>
                      <SelectItem value="browser">Browser</SelectItem>
                    </SelectContent>
                  </Select>
                </div>
              </div>
              
              <Button 
                onClick={generateArchitecture}
                disabled={isGenerating || !userRequest.trim()}
                className="w-full bg-gradient-to-r from-blue-600 to-purple-600 hover:from-blue-700 hover:to-purple-700"
              >
                {isGenerating ? (
                  <>
                    <RefreshCw className="w-4 h-4 mr-2 animate-spin" />
                    Generating Architecture...
                  </>
                ) : (
                  <>
                    <Rocket className="w-4 h-4 mr-2" />
                    Generate AI Architecture
                  </>
                )}
              </Button>
            </CardContent>
          </Card>
          
          {/* Advanced Settings */}
          <Card className="border-slate-700">
            <CardHeader>
              <CardTitle className="text-lg flex items-center gap-2">
                <Settings className="w-5 h-5 text-gray-400" />
                Advanced Settings
              </CardTitle>
            </CardHeader>
            <CardContent className="space-y-3">
              <div>
                <div className="flex justify-between mb-1">
                  <Label className="text-sm">Consciousness Level</Label>
                  <span className="text-xs text-gray-400">{(advancedSettings.consciousness_level * 100).toFixed(0)}%</span>
                </div>
                <input
                  type="range"
                  min="0"
                  max="100"
                  value={advancedSettings.consciousness_level * 100}
                  onChange={(e) => setAdvancedSettings(prev => ({ ...prev, consciousness_level: parseInt(e.target.value) / 100 }))}
                  className="w-full"
                />
              </div>
              
              <div>
                <div className="flex justify-between mb-1">
                  <Label className="text-sm">Creativity Factor</Label>
                  <span className="text-xs text-gray-400">{(advancedSettings.creativity_factor * 100).toFixed(0)}%</span>
                </div>
                <input
                  type="range"
                  min="0"
                  max="100"
                  value={advancedSettings.creativity_factor * 100}
                  onChange={(e) => setAdvancedSettings(prev => ({ ...prev, creativity_factor: parseInt(e.target.value) / 100 }))}
                  className="w-full"
                />
              </div>
              
              <div className="space-y-2">
                <label className="flex items-center gap-2">
                  <input
                    type="checkbox"
                    checked={advancedSettings.include_monitoring}
                    onChange={(e) => setAdvancedSettings(prev => ({ ...prev, include_monitoring: e.target.checked }))}
                    className="rounded"
                  />
                  <span className="text-sm">Include Monitoring</span>
                </label>
                <label className="flex items-center gap-2">
                  <input
                    type="checkbox"
                    checked={advancedSettings.include_testing}
                    onChange={(e) => setAdvancedSettings(prev => ({ ...prev, include_testing: e.target.checked }))}
                    className="rounded"
                  />
                  <span className="text-sm">Include Testing</span>
                </label>
                <label className="flex items-center gap-2">
                  <input
                    type="checkbox"
                    checked={advancedSettings.include_documentation}
                    onChange={(e) => setAdvancedSettings(prev => ({ ...prev, include_documentation: e.target.checked }))}
                    className="rounded"
                  />
                  <span className="text-sm">Include Documentation</span>
                </label>
              </div>
            </CardContent>
          </Card>
        </div>
        
        {/* Right Panel - Results */}
        <div className="lg:col-span-2">
          <Card className="border-slate-700 h-full">
            <CardHeader>
              <div className="flex items-center justify-between">
                <CardTitle className="text-lg">Architecture Results</CardTitle>
                {workflowStatus && (
                  <div className="flex items-center gap-2">
                    {getStatusIcon(workflowStatus.status)}
                    <span className={cn("text-sm font-medium", getStatusColor(workflowStatus.status))}>
                      {workflowStatus.status.replace(/_/g, ' ').toUpperCase()}
                    </span>
                  </div>
                )}
              </div>
            </CardHeader>
            <CardContent>
              <Tabs value={activeTab} onValueChange={setActiveTab}>
                <TabsList className="grid grid-cols-4 w-full">
                  <TabsTrigger value="request">Request</TabsTrigger>
                  <TabsTrigger value="status">Status</TabsTrigger>
                  <TabsTrigger value="architecture">Architecture</TabsTrigger>
                  <TabsTrigger value="code">Code</TabsTrigger>
                </TabsList>
                
                <TabsContent value="request" className="space-y-4">
                  {userRequest ? (
                    <div className="p-4 bg-slate-900/50 rounded-lg">
                      <p className="text-sm text-gray-300">{userRequest}</p>
                      <div className="mt-4 flex gap-2">
                        <Badge variant="outline">{selectedModel}</Badge>
                        <Badge variant="outline">{deploymentTarget}</Badge>
                      </div>
                    </div>
                  ) : (
                    <div className="text-center py-12 text-gray-500">
                      Enter a project request to get started
                    </div>
                  )}
                </TabsContent>
                
                <TabsContent value="status" className="space-y-4">
                  {workflowStatus ? (
                    <div className="space-y-4">
                      {workflowStatus.progress_percentage !== undefined && (
                        <div>
                          <div className="flex justify-between mb-2">
                            <span className="text-sm">Progress</span>
                            <span className="text-sm">{workflowStatus.progress_percentage.toFixed(0)}%</span>
                          </div>
                          <Progress value={workflowStatus.progress_percentage} className="h-2" />
                        </div>
                      )}
                      
                      {workflowStatus.current_stage !== undefined && (
                        <div className="p-3 bg-slate-900/50 rounded-lg">
                          <div className="text-sm text-gray-400">Current Stage</div>
                          <div className="text-lg font-medium">Stage {workflowStatus.current_stage} of 11</div>
                        </div>
                      )}
                      
                      {workflowStatus.error && (
                        <Alert className="border-red-500/50">
                          <AlertCircle className="h-4 w-4" />
                          <AlertDescription>{workflowStatus.error}</AlertDescription>
                        </Alert>
                      )}
                      
                      {workflowStatus.status === 'completed' && workflowStatus.solution && (
                        <div className="p-4 bg-green-900/20 border border-green-500/30 rounded-lg">
                          <div className="flex items-center gap-2 mb-2">
                            <CheckCircle className="w-5 h-5 text-green-400" />
                            <span className="font-medium text-green-400">Architecture Generated Successfully!</span>
                          </div>
                          <p className="text-sm text-gray-300">
                            Your AI architecture is ready. Check the Architecture and Code tabs for details.
                          </p>
                        </div>
                      )}
                    </div>
                  ) : (
                    <div className="text-center py-12 text-gray-500">
                      No active workflow
                    </div>
                  )}
                </TabsContent>
                
                <TabsContent value="architecture" className="space-y-4">
                  {workflowStatus?.solution ? (
                    <div className="space-y-4">
                      {/* Data Plan */}
                      {workflowStatus.solution.data_plan && (
                        <div className="p-4 bg-slate-900/50 rounded-lg space-y-3">
                          <h3 className="font-medium flex items-center gap-2">
                            <Database className="w-4 h-4" />
                            Data Strategy
                          </h3>
                          <div className="grid grid-cols-2 gap-3 text-sm">
                            {workflowStatus.solution.data_plan.selected_datasets?.map((dataset: any, idx: number) => (
                              <div key={idx} className="p-2 bg-slate-800/50 rounded">
                                <div className="font-medium">{dataset.name}</div>
                                <div className="text-xs text-gray-400">{dataset.source}</div>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}
                      
                      {/* Model Architecture */}
                      {workflowStatus.solution.model_architecture && (
                        <div className="p-4 bg-slate-900/50 rounded-lg space-y-3">
                          <h3 className="font-medium flex items-center gap-2">
                            <Cpu className="w-4 h-4" />
                            Model Architecture
                          </h3>
                          <div className="space-y-2 text-sm">
                            <div className="flex justify-between">
                              <span className="text-gray-400">Type:</span>
                              <span>{workflowStatus.solution.model_architecture.model_type}</span>
                            </div>
                            <div className="flex justify-between">
                              <span className="text-gray-400">Framework:</span>
                              <span>{workflowStatus.solution.model_architecture.framework}</span>
                            </div>
                            {workflowStatus.solution.model_architecture.layers && (
                              <div className="mt-3">
                                <div className="text-gray-400 mb-1">Layers:</div>
                                <div className="space-y-1">
                                  {workflowStatus.solution.model_architecture.layers.map((layer: any, idx: number) => (
                                    <div key={idx} className="pl-3 text-xs font-mono">
                                      {layer.type}: {layer.config?.units || layer.config?.filters || ''}
                                    </div>
                                  ))}
                                </div>
                              </div>
                            )}
                          </div>
                        </div>
                      )}
                      
                      {/* Training Config */}
                      {workflowStatus.solution.training_config && (
                        <div className="p-4 bg-slate-900/50 rounded-lg space-y-3">
                          <h3 className="font-medium flex items-center gap-2">
                            <BarChart className="w-4 h-4" />
                            Training Configuration
                          </h3>
                          <div className="grid grid-cols-2 gap-3 text-sm">
                            {Object.entries(workflowStatus.solution.training_config.hyperparameters || {}).map(([key, value]: [string, any]) => (
                              <div key={key} className="flex justify-between">
                                <span className="text-gray-400">{key.replace(/_/g, ' ')}:</span>
                                <span>{value}</span>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}
                      
                      {/* Action Buttons */}
                      <div className="flex gap-3">
                        {workflowStatus.solution.jupyter_notebook && (
                          <Button onClick={downloadNotebook} variant="outline" size="sm">
                            <Download className="w-4 h-4 mr-2" />
                            Download Notebook
                          </Button>
                        )}
                        {workflowStatus.solution.training_pipeline?.training_code && (
                          <Button onClick={exportPythonScript} variant="outline" size="sm">
                            <FileCode className="w-4 h-4 mr-2" />
                            Export Python
                          </Button>
                        )}
                        <Button onClick={deployToCloud} variant="outline" size="sm">
                          <Rocket className="w-4 h-4 mr-2" />
                          Deploy to Cloud
                        </Button>
                      </div>
                    </div>
                  ) : (
                    <div className="text-center py-12 text-gray-500">
                      Architecture will appear here once generated
                    </div>
                  )}
                </TabsContent>
                
                <TabsContent value="code" className="space-y-4">
                  {workflowStatus?.solution?.training_pipeline?.training_code ? (
                    <div className="relative">
                      <div className="absolute top-2 right-2 flex gap-2">
                        <Button
                          onClick={() => copyCode(workflowStatus.solution.training_pipeline.training_code)}
                          variant="ghost"
                          size="sm"
                        >
                          <Copy className="w-4 h-4" />
                        </Button>
                      </div>
                      <pre className="p-4 bg-slate-900 rounded-lg overflow-x-auto">
                        <code className="text-sm text-gray-300">
                          {formatCode(workflowStatus.solution.training_pipeline.training_code)}
                        </code>
                      </pre>
                    </div>
                  ) : (
                    <div className="text-center py-12 text-gray-500">
                      Code will be generated with the architecture
                    </div>
                  )}
                </TabsContent>
              </Tabs>
            </CardContent>
          </Card>
        </div>
      </div>
      
      {/* Error Alert */}
      {error && (
        <Alert className="border-red-500/50">
          <AlertCircle className="h-4 w-4" />
          <AlertDescription>{error}</AlertDescription>
        </Alert>
      )}

      {/* Unified Questionnaire Modal */}
      <UnifiedQuestionnaire
        isOpen={showQuestionnaire}
        onClose={() => setShowQuestionnaire(false)}
        onSubmit={handleQuestionnaireSubmit}
        isSubmitting={isSubmittingQuestionnaire}
      />
    </div>
  );
}