bakery-ia/frontend/src/components/domain/onboarding/steps/MLTrainingStep.tsx

import React, { useState, useEffect, useRef } from 'react';
import { Brain, Activity, Zap, CheckCircle, AlertCircle, TrendingUp, Upload, Database } from 'lucide-react';
import { Button, Card, Badge } from '../../../ui';
import { OnboardingStepProps } from '../OnboardingWizard';
import { useAuthUser } from '../../../../stores/auth.store';
import { useCurrentTenant } from '../../../../stores/tenant.store';
import { useAlertActions } from '../../../../stores/alerts.store';
import { 
  WebSocketService, 
  TrainingProgressMessage, 
  TrainingCompletedMessage, 
  TrainingErrorMessage 
} from '../../../../services/realtime/websocket.service';

interface TrainingMetrics {
  accuracy: number;
  mape: number;
  mae: number;
  rmse: number;
}

interface TrainingLog {
  timestamp: string;
  message: string;
  level: 'info' | 'warning' | 'error' | 'success';
}

interface TrainingJob {
  id: string;
  status: 'pending' | 'running' | 'completed' | 'failed';
  progress: number;
  started_at?: string;
  completed_at?: string;
  error_message?: string;
  metrics?: TrainingMetrics;
}

// Real training service using backend APIs
class TrainingApiService {
  private async apiCall(endpoint: string, options: RequestInit = {}) {
    const response = await fetch(`/api${endpoint}`, {
      ...options,
      headers: {
        'Content-Type': 'application/json',
        ...options.headers,
      },
    });

    if (!response.ok) {
      throw new Error(`API call failed: ${response.statusText}`);
    }

    return response.json();
  }

  async startTrainingJob(tenantId: string, startDate?: string, endDate?: string): Promise<TrainingJob> {
    return this.apiCall(`/tenants/${tenantId}/training/jobs`, {
      method: 'POST',
      body: JSON.stringify({
        start_date: startDate,
        end_date: endDate
      }),
    });
  }

  async getTrainingJob(tenantId: string, jobId: string): Promise<TrainingJob> {
    return this.apiCall(`/tenants/${tenantId}/training/jobs/${jobId}`);
  }

  async getTrainingJobs(tenantId: string): Promise<TrainingJob[]> {
    return this.apiCall(`/tenants/${tenantId}/training/jobs`);
  }
}

const trainingService = new TrainingApiService();

export const MLTrainingStep: React.FC<OnboardingStepProps> = ({
  data,
  onDataChange,
  onNext,
  onPrevious,
  isFirstStep,
  isLastStep
}) => {
  const user = useAuthUser();
  const currentTenant = useCurrentTenant();
  const { createAlert } = useAlertActions();
  
  const [trainingStatus, setTrainingStatus] = useState<'idle' | 'validating' | 'training' | 'completed' | 'failed'>(
    data.trainingStatus || 'idle'
  );
  const [progress, setProgress] = useState(data.trainingProgress || 0);
  const [currentJob, setCurrentJob] = useState<TrainingJob | null>(data.trainingJob || null);
  const [trainingLogs, setTrainingLogs] = useState<TrainingLog[]>(data.trainingLogs || []);
  const [metrics, setMetrics] = useState<TrainingMetrics | null>(data.trainingMetrics || null);
  const [currentStep, setCurrentStep] = useState<string>('');
  const [estimatedTimeRemaining, setEstimatedTimeRemaining] = useState<number>(0);
  
  const wsRef = useRef<WebSocketService | null>(null);

  // Validate that required data is available for training
  const validateDataRequirements = (): { isValid: boolean; missingItems: string[] } => {
    const missingItems: string[] = [];

    console.log('MLTrainingStep - Validating data requirements');
    console.log('MLTrainingStep - Current data:', data);
    console.log('MLTrainingStep - allStepData keys:', Object.keys(data.allStepData || {}));

    // Get data from previous steps
    const dataProcessingData = data.allStepData?.['data-processing'];
    const reviewData = data.allStepData?.['review'];
    const inventoryData = data.allStepData?.['inventory'];

    console.log('MLTrainingStep - dataProcessingData:', dataProcessingData);
    console.log('MLTrainingStep - reviewData:', reviewData);
    console.log('MLTrainingStep - inventoryData:', inventoryData);

    // Check if sales data was processed
    const hasProcessingResults = dataProcessingData?.processingResults && 
                                dataProcessingData.processingResults.is_valid &&
                                dataProcessingData.processingResults.total_records > 0;
    
    if (!hasProcessingResults) {
      missingItems.push('Datos de ventas importados');
    }

    // Check if products were approved in review step
    const hasApprovedProducts = reviewData?.approvedProducts && 
                               reviewData.approvedProducts.length > 0 &&
                               reviewData.reviewCompleted;
    
    if (!hasApprovedProducts) {
      missingItems.push('Productos aprobados en revisión');
    }

    // Check if inventory was configured
    const hasInventoryConfig = inventoryData?.inventoryConfigured &&
                              inventoryData?.inventoryItems &&
                              inventoryData.inventoryItems.length > 0;
    
    if (!hasInventoryConfig) {
      missingItems.push('Inventario configurado');
    }

    // Check if we have enough data for training
    if (dataProcessingData?.processingResults?.total_records && 
        dataProcessingData.processingResults.total_records < 10) {
      missingItems.push('Suficientes registros de ventas (mínimo 10)');
    }

    console.log('MLTrainingStep - Validation result:', {
      isValid: missingItems.length === 0,
      missingItems,
      hasProcessingResults,
      hasApprovedProducts,
      hasInventoryConfig
    });

    return {
      isValid: missingItems.length === 0,
      missingItems
    };
  };

  const addLog = (message: string, level: TrainingLog['level'] = 'info') => {
    const newLog: TrainingLog = {
      timestamp: new Date().toISOString(),
      message,
      level
    };
    setTrainingLogs(prev => [...prev, newLog]);
  };

  const startTraining = async () => {
    const tenantId = currentTenant?.id || user?.tenant_id;
    if (!tenantId) {
      createAlert({
        type: 'error',
        category: 'system',
        priority: 'high',
        title: 'Error',
        message: 'No se pudo obtener información del tenant',
        source: 'onboarding'
      });
      return;
    }

    // Validate data requirements
    const validation = validateDataRequirements();
    if (!validation.isValid) {
      createAlert({
        type: 'error',
        category: 'system',
        priority: 'high',
        title: 'Datos insuficientes para entrenamiento',
        message: `Faltan los siguientes elementos: ${validation.missingItems.join(', ')}`,
        source: 'onboarding'
      });
      return;
    }

    setTrainingStatus('validating');
    addLog('Validando disponibilidad de datos...', 'info');

    try {
      // Start training job
      addLog('Iniciando trabajo de entrenamiento ML...', 'info');
      const job = await trainingService.startTrainingJob(tenantId);
      
      setCurrentJob(job);
      setTrainingStatus('training');
      addLog(`Trabajo de entrenamiento iniciado: ${job.id}`, 'success');

      // Initialize WebSocket connection for real-time updates
      const ws = new WebSocketService(tenantId, job.id);
      wsRef.current = ws;

      // Set up WebSocket event listeners
      ws.subscribe('progress', (message: TrainingProgressMessage) => {
        console.log('Training progress received:', message);
        setProgress(message.progress.percentage);
        setCurrentStep(message.progress.current_step);
        setEstimatedTimeRemaining(message.progress.estimated_time_remaining);
        
        addLog(
          `${message.progress.current_step} - ${message.progress.products_completed}/${message.progress.products_total} productos procesados (${message.progress.percentage}%)`,
          'info'
        );
      });

      ws.subscribe('completed', (message: TrainingCompletedMessage) => {
        console.log('Training completed:', message);
        setTrainingStatus('completed');
        setProgress(100);
        
        const metrics: TrainingMetrics = {
          accuracy: message.results.performance_metrics.accuracy,
          mape: message.results.performance_metrics.mape,
          mae: message.results.performance_metrics.mae,
          rmse: message.results.performance_metrics.rmse
        };
        
        setMetrics(metrics);
        addLog('¡Entrenamiento ML completado exitosamente!', 'success');
        addLog(`${message.results.successful_trainings} modelos creados exitosamente`, 'success');
        addLog(`Duración total: ${Math.round(message.results.training_duration / 60)} minutos`, 'info');
        
        createAlert({
          type: 'success',
          category: 'system',
          priority: 'medium',
          title: 'Entrenamiento completado',
          message: `Tu modelo de IA ha sido entrenado exitosamente. Precisión: ${(metrics.accuracy * 100).toFixed(1)}%`,
          source: 'onboarding'
        });

        // Update parent data
        onDataChange({
          ...data,
          trainingStatus: 'completed',
          trainingProgress: 100,
          trainingJob: { ...job, status: 'completed', progress: 100, metrics },
          trainingLogs,
          trainingMetrics: metrics
        });

        // Disconnect WebSocket
        ws.disconnect();
        wsRef.current = null;
      });

      ws.subscribe('error', (message: TrainingErrorMessage) => {
        console.error('Training error received:', message);
        setTrainingStatus('failed');
        addLog(`Error en entrenamiento: ${message.error}`, 'error');
        
        createAlert({
          type: 'error',
          category: 'system',
          priority: 'high',
          title: 'Error en entrenamiento',
          message: message.error,
          source: 'onboarding'
        });

        // Disconnect WebSocket
        ws.disconnect();
        wsRef.current = null;
      });

      // Connect to WebSocket
      await ws.connect();
      addLog('Conectado a WebSocket para actualizaciones en tiempo real', 'info');

    } catch (error) {
      console.error('Training start error:', error);
      setTrainingStatus('failed');
      const errorMessage = error instanceof Error ? error.message : 'Error al iniciar entrenamiento';
      addLog(`Error: ${errorMessage}`, 'error');
      
      createAlert({
        type: 'error',
        category: 'system',
        priority: 'high',
        title: 'Error al iniciar entrenamiento',
        message: errorMessage,
        source: 'onboarding'
      });

      // Clean up WebSocket if it was created
      if (wsRef.current) {
        wsRef.current.disconnect();
        wsRef.current = null;
      }
    }
  };

  // Cleanup WebSocket on unmount
  useEffect(() => {
    return () => {
      if (wsRef.current) {
        wsRef.current.disconnect();
        wsRef.current = null;
      }
    };
  }, []);

  useEffect(() => {
    // Auto-start training if all requirements are met and not already started
    const validation = validateDataRequirements();
    console.log('MLTrainingStep - useEffect validation:', validation);
    
    if (validation.isValid && trainingStatus === 'idle' && data.autoStartTraining) {
      console.log('MLTrainingStep - Auto-starting training...');
      // Auto-start after a brief delay to allow user to see the step
      const timer = setTimeout(() => {
        startTraining();
      }, 1000);
      return () => clearTimeout(timer);
    }
  }, [data.allStepData, data.autoStartTraining, trainingStatus]);

  const getStatusIcon = () => {
    switch (trainingStatus) {
      case 'idle': return <Brain className="w-8 h-8 text-[var(--color-primary)]" />;
      case 'validating': return <Database className="w-8 h-8 text-[var(--color-info)] animate-pulse" />;
      case 'training': return <Activity className="w-8 h-8 text-[var(--color-info)] animate-pulse" />;
      case 'completed': return <CheckCircle className="w-8 h-8 text-[var(--color-success)]" />;
      case 'failed': return <AlertCircle className="w-8 h-8 text-[var(--color-error)]" />;
      default: return <Brain className="w-8 h-8 text-[var(--color-primary)]" />;
    }
  };

  const getStatusColor = () => {
    switch (trainingStatus) {
      case 'completed': return 'text-[var(--color-success)]';
      case 'failed': return 'text-[var(--color-error)]';
      case 'training': 
      case 'validating': return 'text-[var(--color-info)]';
      default: return 'text-[var(--text-primary)]';
    }
  };

  const getStatusMessage = () => {
    switch (trainingStatus) {
      case 'idle': return 'Listo para entrenar tu asistente IA';
      case 'validating': return 'Validando datos para entrenamiento...';
      case 'training': return 'Entrenando modelo de predicción...';
      case 'completed': return '¡Tu asistente IA está listo!';
      case 'failed': return 'Error en el entrenamiento';
      default: return 'Estado desconocido';
    }
  };

  // Check data requirements for display
  const validation = validateDataRequirements();

  if (!validation.isValid) {
    return (
      <div className="space-y-8">
        <div className="text-center py-16">
          <Upload className="w-16 h-16 text-gray-400 mx-auto mb-4" />
          <h3 className="text-2xl font-bold text-gray-600 mb-4">
            Datos insuficientes para entrenamiento
          </h3>
          <p className="text-gray-500 mb-6 max-w-md mx-auto">
            Para entrenar tu modelo de IA, necesitamos que completes los siguientes elementos:
          </p>
          
          <Card className="p-6 max-w-md mx-auto mb-6">
            <h4 className="font-semibold mb-4 text-[var(--text-primary)]">Elementos requeridos:</h4>
            <ul className="space-y-2 text-left">
              {validation.missingItems.map((item, index) => (
                <li key={index} className="flex items-center space-x-2">
                  <AlertCircle className="w-4 h-4 text-red-500" />
                  <span className="text-sm text-[var(--text-secondary)]">{item}</span>
                </li>
              ))}
            </ul>
          </Card>

          <p className="text-sm text-[var(--text-secondary)] mb-4">
            Una vez completados estos elementos, el entrenamiento se iniciará automáticamente.
          </p>
          
          <Button 
            variant="outline" 
            onClick={onPrevious}
            disabled={isFirstStep}
          >
            Volver al paso anterior
          </Button>
        </div>
      </div>
    );
  }

  return (
    <div className="space-y-8">
      {/* Header */}
      <div className="text-center">
        <div className="w-20 h-20 bg-gradient-to-br from-[var(--color-primary)] to-[var(--color-secondary)] rounded-full flex items-center justify-center mx-auto mb-6">
          {getStatusIcon()}
        </div>
        <h2 className={`text-3xl font-bold mb-4 ${getStatusColor()}`}>
          Entrenamiento de IA
        </h2>
        <p className={`text-lg mb-2 ${getStatusColor()}`}>
          {getStatusMessage()}
        </p>
        <p className="text-[var(--text-secondary)]">
          Creando tu asistente inteligente personalizado con tus datos de ventas e inventario
        </p>
      </div>

      {/* Progress Bar */}
      <Card className="p-6">
        <div className="flex justify-between items-center mb-2">
          <span className="text-sm font-medium text-[var(--text-primary)]">Progreso del entrenamiento</span>
          <span className="text-sm text-[var(--text-secondary)]">{progress.toFixed(1)}%</span>
        </div>
        
        {currentStep && (
          <div className="mb-2">
            <span className="text-xs text-[var(--text-secondary)]">Paso actual: </span>
            <span className="text-xs font-medium text-[var(--text-primary)]">{currentStep}</span>
          </div>
        )}
        
        {estimatedTimeRemaining > 0 && (
          <div className="mb-4">
            <span className="text-xs text-[var(--text-secondary)]">Tiempo estimado restante: </span>
            <span className="text-xs font-medium text-[var(--color-primary)]">
              {Math.round(estimatedTimeRemaining / 60)} minutos
            </span>
          </div>
        )}
        
        <div className="w-full bg-[var(--bg-secondary)] rounded-full h-3">
          <div 
            className="bg-gradient-to-r from-[var(--color-primary)] to-[var(--color-secondary)] h-3 rounded-full transition-all duration-300"
            style={{ width: `${progress}%` }}
          />
        </div>
      </Card>

      {/* Training Logs */}
      <Card className="p-6">
        <h3 className="text-lg font-semibold mb-4 flex items-center">
          <Activity className="w-5 h-5 mr-2" />
          Registro de entrenamiento
        </h3>
        <div className="space-y-2 max-h-64 overflow-y-auto">
          {trainingLogs.length === 0 ? (
            <p className="text-[var(--text-secondary)] italic">Esperando inicio de entrenamiento...</p>
          ) : (
            trainingLogs.map((log, index) => (
              <div key={index} className="flex items-start space-x-3 p-2 rounded">
                <div className={`w-2 h-2 rounded-full mt-2 ${
                  log.level === 'success' ? 'bg-green-500' :
                  log.level === 'error' ? 'bg-red-500' :
                  log.level === 'warning' ? 'bg-yellow-500' :
                  'bg-blue-500'
                }`} />
                <div className="flex-1">
                  <p className="text-sm text-[var(--text-primary)]">{log.message}</p>
                  <p className="text-xs text-[var(--text-secondary)]">
                    {new Date(log.timestamp).toLocaleTimeString()}
                  </p>
                </div>
              </div>
            ))
          )}
        </div>
      </Card>

      {/* Training Metrics */}
      {metrics && trainingStatus === 'completed' && (
        <Card className="p-6">
          <h3 className="text-lg font-semibold mb-4 flex items-center">
            <TrendingUp className="w-5 h-5 mr-2" />
            Métricas del modelo
          </h3>
          <div className="grid grid-cols-2 md:grid-cols-4 gap-4">
            <div className="text-center p-3 bg-[var(--bg-secondary)] rounded-lg">
              <p className="text-2xl font-bold text-[var(--color-success)]">
                {(metrics.accuracy * 100).toFixed(1)}%
              </p>
              <p className="text-xs text-[var(--text-secondary)]">Precisión</p>
            </div>
            <div className="text-center p-3 bg-[var(--bg-secondary)] rounded-lg">
              <p className="text-2xl font-bold text-[var(--color-info)]">
                {metrics.mape.toFixed(1)}%
              </p>
              <p className="text-xs text-[var(--text-secondary)]">MAPE</p>
            </div>
            <div className="text-center p-3 bg-[var(--bg-secondary)] rounded-lg">
              <p className="text-2xl font-bold text-[var(--color-primary)]">
                {metrics.mae.toFixed(2)}
              </p>
              <p className="text-xs text-[var(--text-secondary)]">MAE</p>
            </div>
            <div className="text-center p-3 bg-[var(--bg-secondary)] rounded-lg">
              <p className="text-2xl font-bold text-[var(--color-secondary)]">
                {metrics.rmse.toFixed(2)}
              </p>
              <p className="text-xs text-[var(--text-secondary)]">RMSE</p>
            </div>
          </div>
        </Card>
      )}

      {/* Manual Start Button (if not auto-started) */}
      {trainingStatus === 'idle' && (
        <Card className="p-6 text-center">
          <Button 
            onClick={startTraining}
            className="bg-[var(--color-primary)] hover:bg-[var(--color-primary)]/90"
            size="lg"
          >
            <Zap className="w-5 h-5 mr-2" />
            Iniciar entrenamiento
          </Button>
        </Card>
      )}

      {/* Navigation */}
      <div className="flex justify-between pt-6 border-t border-[var(--border-primary)]">
        <Button 
          variant="outline" 
          onClick={onPrevious}
          disabled={isFirstStep}
        >
          Anterior
        </Button>
        <Button 
          onClick={onNext}
          disabled={trainingStatus !== 'completed'}
          className="bg-[var(--color-primary)] hover:bg-[var(--color-primary)]/90"
        >
          {isLastStep ? 'Finalizar' : 'Siguiente'}
        </Button>
      </div>
    </div>
  );
};