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

import React, { useState, useCallback, useEffect } from 'react';
import { Button } from '../../../ui/Button';
import { useCurrentTenant } from '../../../../stores/tenant.store';
import { useCreateTrainingJob, useTrainingWebSocket, useTrainingJobStatus } from '../../../../api/hooks/training';

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;
}

export const MLTrainingStep: React.FC<MLTrainingStepProps> = ({
  onComplete
}) => {
  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 currentTenant = useCurrentTenant();
  const createTrainingJob = useCreateTrainingJob();

  // 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
    });
  }, []);

  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 HTTP polling (fallback only)
  useEffect(() => {
    if (!jobStatus || !jobId || trainingProgress?.stage === 'completed') {
      return;
    }

    console.log('📊 HTTP fallback status update:', jobStatus);

    // Check if training completed via HTTP polling fallback
    if (jobStatus.status === 'completed' && trainingProgress?.stage !== 'completed') {
      console.log('✅ Training completion detected via HTTP fallback');
      setTrainingProgress({
        stage: 'completed',
        progress: 100,
        message: '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 via HTTP fallback');
      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 from HTTP polling fallback
      const currentProgress = trainingProgress?.progress || 0;
      if (jobStatus.progress > currentProgress) {
        console.log(`📈 Progress update via HTTP fallback: ${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]);

  // 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);
      
      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 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`;
    }
  };

  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-2">
                    <div className="w-full bg-[var(--bg-tertiary)] rounded-full h-2">
                      <div
                        className="bg-[var(--color-primary)] h-2 rounded-full transition-all duration-300"
                        style={{ width: `${trainingProgress.progress}%` }}
                      />
                    </div>

                    <div className="flex justify-between text-xs text-[var(--text-tertiary)]">
                      <span>{trainingProgress.currentStep || 'Procesando...'}</span>
                      <div className="flex items-center gap-2">
                        {jobId && (
                          <span className={`text-xs ${isConnected ? 'text-green-500' : 'text-red-500'}`}>
                            {isConnected ? '🟢 Conectado' : '🔴 Desconectado'}
                          </span>
                        )}
                        {trainingProgress.estimatedTimeRemaining && (
                          <span>Tiempo estimado: {formatTime(trainingProgress.estimatedTimeRemaining)}</span>
                        )}
                      </div>
                    </div>
                  </div>
                )}
              </div>
            </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>
  );
};