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

import React, { useState, useCallback, useEffect } from 'react';
import { useNavigate } from 'react-router-dom';
import { useTranslation } from 'react-i18next';
import { Button } from '../../../ui/Button';
import { useCurrentTenant } from '../../../../stores/tenant.store';
import { useCreateTrainingJob, useTrainingWebSocket, useTrainingJobStatus } from '../../../../api/hooks/training';
import { Info } from 'lucide-react';

interface MLTrainingStepProps {
  onNext: () => void;
  onPrevious: () => void;
  onComplete: (data?: any) => void;
  isFirstStep: boolean;
  isLastStep: boolean;
}

interface TrainingProgress {
  stage: string;
  progress: number;
  message: string;
  currentStep?: string;
  estimatedTimeRemaining?: number;
  estimatedCompletionTime?: string;
}

export const MLTrainingStep: React.FC<MLTrainingStepProps> = ({
  onComplete
}) => {
  const { t } = useTranslation();
  const navigate = useNavigate();
  const [trainingProgress, setTrainingProgress] = useState<TrainingProgress | null>(null);
  const [isTraining, setIsTraining] = useState(false);
  const [error, setError] = useState<string>('');
  const [jobId, setJobId] = useState<string | null>(null);
  const [trainingStartTime, setTrainingStartTime] = useState<number | null>(null);
  const [showSkipOption, setShowSkipOption] = useState(false);

  const currentTenant = useCurrentTenant();
  const createTrainingJob = useCreateTrainingJob();

  // Check if training has been running for more than 2 minutes
  useEffect(() => {
    if (trainingStartTime && isTraining && !showSkipOption) {
      const checkTimer = setInterval(() => {
        const elapsedTime = (Date.now() - trainingStartTime) / 1000; // in seconds
        if (elapsedTime > 120) { // 2 minutes
          setShowSkipOption(true);
          clearInterval(checkTimer);
        }
      }, 5000); // Check every 5 seconds

      return () => clearInterval(checkTimer);
    }
  }, [trainingStartTime, isTraining, showSkipOption]);

  // Memoized WebSocket callbacks to prevent reconnections
  const handleProgress = useCallback((data: any) => {
    setTrainingProgress({
      stage: 'training',
      progress: data.data?.progress || 0,
      message: data.data?.message || 'Entrenando modelo...',
      currentStep: data.data?.current_step,
      estimatedTimeRemaining: data.data?.estimated_time_remaining_seconds || data.data?.estimated_time_remaining,
      estimatedCompletionTime: data.data?.estimated_completion_time
    });
  }, []);

  const handleCompleted = useCallback((_data: any) => {
    setTrainingProgress({
      stage: 'completed',
      progress: 100,
      message: 'Entrenamiento completado exitosamente'
    });
    setIsTraining(false);
    
    setTimeout(() => {
      onComplete({
        jobId: jobId,
        success: true,
        message: 'Modelo entrenado correctamente'
      });
    }, 2000);
  }, [onComplete, jobId]);

  const handleError = useCallback((data: any) => {
    setError(data.data?.error || data.error || 'Error durante el entrenamiento');
    setIsTraining(false);
    setTrainingProgress(null);
  }, []);

  const handleStarted = useCallback((_data: any) => {
    setTrainingProgress({
      stage: 'starting',
      progress: 5,
      message: 'Iniciando entrenamiento del modelo...'
    });
  }, []);

  // WebSocket for real-time training progress - only connect when we have a jobId
  const { isConnected, connectionError } = useTrainingWebSocket(
    currentTenant?.id || '',
    jobId || '',
    undefined, // token will be handled by the service
    jobId ? {
      onProgress: handleProgress,
      onCompleted: handleCompleted,
      onError: handleError,
      onStarted: handleStarted
    } : undefined
  );

  // Smart fallback polling - automatically disabled when WebSocket is connected
  const { data: jobStatus } = useTrainingJobStatus(
    currentTenant?.id || '',
    jobId || '',
    {
      enabled: !!jobId && !!currentTenant?.id,
      isWebSocketConnected: isConnected, // This will disable HTTP polling when WebSocket is connected
    }
  );

  // Handle training status updates from React Query cache (updated by WebSocket or HTTP fallback)
  useEffect(() => {
    if (!jobStatus || !jobId || trainingProgress?.stage === 'completed') {
      return;
    }

    console.log('📊 Training status update from cache:', jobStatus,
                `(source: ${isConnected ? 'WebSocket' : 'HTTP polling'})`);

    // Check if training completed
    if (jobStatus.status === 'completed' && trainingProgress?.stage !== 'completed') {
      console.log(`✅ Training completion detected (source: ${isConnected ? 'WebSocket' : 'HTTP polling'})`);
      setTrainingProgress({
        stage: 'completed',
        progress: 100,
        message: isConnected
          ? 'Entrenamiento completado exitosamente'
          : 'Entrenamiento completado exitosamente (detectado por verificación HTTP)'
      });
      setIsTraining(false);

      setTimeout(() => {
        onComplete({
          jobId: jobId,
          success: true,
          message: 'Modelo entrenado correctamente',
          detectedViaPolling: true
        });
      }, 2000);
    } else if (jobStatus.status === 'failed') {
      console.log(`❌ Training failure detected (source: ${isConnected ? 'WebSocket' : 'HTTP polling'})`);
      setError('Error detectado durante el entrenamiento (verificación de estado)');
      setIsTraining(false);
      setTrainingProgress(null);
    } else if (jobStatus.status === 'running' && jobStatus.progress !== undefined) {
      // Update progress if we have newer information
      const currentProgress = trainingProgress?.progress || 0;
      if (jobStatus.progress > currentProgress) {
        console.log(`📈 Progress update (source: ${isConnected ? 'WebSocket' : 'HTTP polling'}): ${jobStatus.progress}%`);
        setTrainingProgress(prev => ({
          ...prev,
          stage: 'training',
          progress: jobStatus.progress,
          message: jobStatus.message || 'Entrenando modelo...',
          currentStep: jobStatus.current_step
        }) as TrainingProgress);
      }
    }
  }, [jobStatus, jobId, trainingProgress?.stage, onComplete, isConnected]);

  // Auto-trigger training when component mounts
  useEffect(() => {
    if (currentTenant?.id && !isTraining && !trainingProgress && !error) {
      console.log('🚀 Auto-starting ML training for tenant:', currentTenant.id);
      handleStartTraining();
    }
  }, [currentTenant?.id]); // Only run when tenant is available

  const handleStartTraining = async () => {
    if (!currentTenant?.id) {
      setError('No se encontró información del tenant');
      return;
    }

    setIsTraining(true);
    setError('');
    setTrainingProgress({
      stage: 'preparing',
      progress: 0,
      message: 'Preparando datos para entrenamiento...'
    });

    try {
      const response = await createTrainingJob.mutateAsync({
        tenantId: currentTenant.id,
        request: {
          // Use the exact backend schema - all fields are optional
          // This will train on all available data
        }
      });

      setJobId(response.job_id);
      setTrainingStartTime(Date.now()); // Track when training started

      setTrainingProgress({
        stage: 'queued',
        progress: 10,
        message: 'Trabajo de entrenamiento en cola...'
      });
    } catch (err) {
      setError('Error al iniciar el entrenamiento del modelo');
      setIsTraining(false);
      setTrainingProgress(null);
    }
  };

  const handleSkipToDashboard = () => {
    // Navigate to dashboard while training continues in background
    console.log('🚀 User chose to skip to dashboard while training continues');
    navigate('/app/dashboard');
  };

  const formatTime = (seconds?: number) => {
    if (!seconds) return '';

    if (seconds < 60) {
      return `${Math.round(seconds)}s`;
    } else if (seconds < 3600) {
      return `${Math.round(seconds / 60)}m`;
    } else {
      return `${Math.round(seconds / 3600)}h ${Math.round((seconds % 3600) / 60)}m`;
    }
  };

  const formatEstimatedCompletionTime = (isoString?: string) => {
    if (!isoString) return '';

    try {
      const completionDate = new Date(isoString);
      const now = new Date();

      // If completion is today, show time only
      if (completionDate.toDateString() === now.toDateString()) {
        return completionDate.toLocaleTimeString('es-ES', {
          hour: '2-digit',
          minute: '2-digit'
        });
      }

      // If completion is another day, show date and time
      return completionDate.toLocaleString('es-ES', {
        month: 'short',
        day: 'numeric',
        hour: '2-digit',
        minute: '2-digit'
      });
    } catch (error) {
      return '';
    }
  };

  return (
    <div className="space-y-6">
      <div className="text-center">
        <p className="text-[var(--text-secondary)] mb-6">
          Perfecto! Ahora entrenaremos automáticamente tu modelo de inteligencia artificial utilizando los datos de ventas 
          e inventario que has proporcionado. Este proceso puede tomar varios minutos.
        </p>
      </div>

      {/* Training Status Card */}
      <div className="bg-[var(--bg-secondary)] rounded-lg p-6">
        <div className="text-center">
          {!isTraining && !trainingProgress && (
            <div className="space-y-4">
              <div className="mx-auto w-16 h-16 bg-[var(--color-primary)]/10 rounded-full flex items-center justify-center">
                <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-[var(--color-primary)]"></div>
              </div>
              <div>
                <h3 className="text-lg font-semibold mb-2">Iniciando Entrenamiento Automático</h3>
                <p className="text-[var(--text-secondary)] text-sm">
                  Preparando el entrenamiento de tu modelo con los datos proporcionados...
                </p>
              </div>
            </div>
          )}

          {trainingProgress && (
            <div className="space-y-4">
              <div className="mx-auto w-16 h-16 relative">
                {trainingProgress.stage === 'completed' ? (
                  <div className="w-16 h-16 bg-[var(--color-success)]/10 rounded-full flex items-center justify-center">
                    <svg className="w-8 h-8 text-[var(--color-success)]" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                      <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M5 13l4 4L19 7" />
                    </svg>
                  </div>
                ) : (
                  <div className="w-16 h-16 bg-[var(--color-primary)]/10 rounded-full flex items-center justify-center">
                    <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-[var(--color-primary)]"></div>
                  </div>
                )}
                
                {trainingProgress.progress > 0 && trainingProgress.stage !== 'completed' && (
                  <div className="absolute -bottom-2 left-1/2 transform -translate-x-1/2">
                    <span className="text-xs font-medium text-[var(--text-tertiary)]">
                      {trainingProgress.progress}%
                    </span>
                  </div>
                )}
              </div>

              <div>
                <h3 className="text-lg font-semibold mb-2">
                  {trainingProgress.stage === 'completed' 
                    ? '¡Entrenamiento Completo!' 
                    : 'Entrenando Modelo IA'
                  }
                </h3>
                <p className="text-[var(--text-secondary)] text-sm mb-4">
                  {trainingProgress.message}
                </p>

                {trainingProgress.stage !== 'completed' && (
                  <div className="space-y-3">
                    {/* Enhanced Progress Bar */}
                    <div className="relative">
                      <div className="w-full bg-[var(--bg-tertiary)] rounded-full h-3 overflow-hidden">
                        <div
                          className="bg-gradient-to-r from-[var(--color-primary)] to-[var(--color-primary)]/80 h-3 rounded-full transition-all duration-500 ease-out relative"
                          style={{ width: `${trainingProgress.progress}%` }}
                        >
                          {/* Animated shimmer effect */}
                          <div className="absolute inset-0 bg-gradient-to-r from-transparent via-white/20 to-transparent animate-shimmer"></div>
                        </div>
                      </div>
                      {/* Progress percentage badge */}
                      <div className="absolute -top-1 left-1/2 transform -translate-x-1/2 -translate-y-full mb-1">
                        <span className="text-xs font-semibold text-[var(--color-primary)] bg-[var(--bg-primary)] px-2 py-1 rounded-full shadow-sm border border-[var(--color-primary)]/20">
                          {trainingProgress.progress}%
                        </span>
                      </div>
                    </div>

                    {/* Training Information */}
                    <div className="flex flex-col gap-2 text-xs text-[var(--text-tertiary)]">
                      {/* Current Step */}
                      <div className="flex justify-between items-center">
                        <span className="font-medium">{trainingProgress.currentStep || t('onboarding:steps.ml_training.progress.data_preparation', 'Procesando...')}</span>
                        {jobId && (
                          <span className={`text-xs px-2 py-0.5 rounded-full ${isConnected ? 'bg-green-100 text-green-700 dark:bg-green-900/30 dark:text-green-400' : 'bg-red-100 text-red-700 dark:bg-red-900/30 dark:text-red-400'}`}>
                            {isConnected ? '● En vivo' : '● Reconectando...'}
                          </span>
                        )}
                      </div>

                      {/* Time Information */}
                      <div className="flex flex-wrap gap-x-4 gap-y-1 text-xs">
                        {trainingProgress.estimatedTimeRemaining && (
                          <div className="flex items-center gap-1">
                            <svg className="w-3.5 h-3.5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M12 8v4l3 3m6-3a9 9 0 11-18 0 9 9 0 0118 0z" />
                            </svg>
                            <span>
                              {t('onboarding:steps.ml_training.estimated_time_remaining', 'Tiempo restante: {{time}}', {
                                time: formatTime(trainingProgress.estimatedTimeRemaining)
                              })}
                            </span>
                          </div>
                        )}
                        {trainingProgress.estimatedCompletionTime && (
                          <div className="flex items-center gap-1">
                            <svg className="w-3.5 h-3.5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M8 7V3m8 4V3m-9 8h10M5 21h14a2 2 0 002-2V7a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z" />
                            </svg>
                            <span>
                              Finalizará: {formatEstimatedCompletionTime(trainingProgress.estimatedCompletionTime)}
                            </span>
                          </div>
                        )}
                      </div>
                    </div>
                  </div>
                )}
              </div>
            </div>
          )}
        </div>
      </div>

      {/* Skip to Dashboard Option - Show after 2 minutes */}
      {showSkipOption && isTraining && trainingProgress?.stage !== 'completed' && (
        <div className="bg-blue-50 dark:bg-blue-900/20 border border-blue-200 dark:border-blue-800 rounded-lg p-4">
          <div className="flex items-start gap-3">
            <div className="flex-shrink-0 mt-0.5">
              <Info className="w-5 h-5 text-blue-600 dark:text-blue-400" />
            </div>
            <div className="flex-1">
              <h4 className="font-medium text-blue-900 dark:text-blue-100 mb-1">
                {t('onboarding:steps.ml_training.skip_to_dashboard.title', '¿Toma demasiado tiempo?')}
              </h4>
              <p className="text-sm text-blue-800 dark:text-blue-200 mb-3">
                {t('onboarding:steps.ml_training.skip_to_dashboard.info', 'El entrenamiento está tardando más de lo esperado. No te preocupes, puedes explorar tu dashboard mientras el modelo termina de entrenarse en segundo plano.')}
              </p>
              <Button
                onClick={handleSkipToDashboard}
                variant="secondary"
                size="sm"
              >
                {t('onboarding:steps.ml_training.skip_to_dashboard.button', 'Ir al Dashboard')}
              </Button>
              <p className="text-xs text-blue-700 dark:text-blue-300 mt-2">
                {t('onboarding:steps.ml_training.skip_to_dashboard.training_continues', 'El entrenamiento continúa en segundo plano')}
              </p>
            </div>
          </div>
        </div>
      )}

      {/* Training Info */}
      <div className="bg-[var(--bg-secondary)] rounded-lg p-4">
        <h4 className="font-medium mb-2">¿Qué sucede durante el entrenamiento?</h4>
        <ul className="text-sm text-[var(--text-secondary)] space-y-1">
          <li>• Análisis de patrones de ventas históricos</li>
          <li>• Creación de modelos predictivos de demanda</li>
          <li>• Optimización de algoritmos de inventario</li>
          <li>• Validación y ajuste de precisión</li>
        </ul>
      </div>

      {(error || connectionError) && (
        <div className="bg-[var(--color-error)]/10 border border-[var(--color-error)]/20 rounded-lg p-4">
          <p className="text-[var(--color-error)]">{error || connectionError}</p>
        </div>
      )}

      {/* Auto-completion when training finishes - no manual buttons needed */}
    </div>
  );
};