// src/hooks/useTrainingProgress.ts
import { useState, useEffect } from 'react';
import { useWebSocket } from '../hooks/useWebSocket';

export interface TrainingProgress {
  job_id: string;
  status: 'pending' | 'running' | 'completed' | 'failed';
  progress: number;
  current_step: string;
  total_steps: number;
  estimated_time_remaining?: number;
  metrics?: Record<string, any>;
}

export interface TrainingProgressUpdate {
  type: 'training_progress' | 'training_completed' | 'training_error';
  job_id: string;
  progress?: TrainingProgress;
  results?: any;
  error?: string;
}

export const useTrainingProgress = (jobId: string | null) => {
  const [progress, setProgress] = useState<TrainingProgress | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [isComplete, setIsComplete] = useState(false);

  const handleMessage = (data: TrainingProgressUpdate) => {
    switch (data.type) {
      case 'training_progress':
        setProgress(data.progress!);
        setError(null);
        break;
      
      case 'training_completed':
        setProgress(prev => ({
          ...prev!,
          status: 'completed',
          progress: 100
        }));
        setIsComplete(true);
        break;
      
      case 'training_error':
        setError(data.error || 'Training failed');
        setProgress(prev => prev ? { ...prev, status: 'failed' } : null);
        break;
    }
  };

  const { isConnected } = useWebSocket({
    endpoint: jobId ? `/training/progress/${jobId}` : '',
    onMessage: handleMessage,
    onError: () => setError('Connection lost'),
    autoConnect: !!jobId
  });

  // Fetch initial status when job ID changes
  useEffect(() => {
    if (jobId) {
      fetchTrainingStatus(jobId);
    }
  }, [jobId]);

  const fetchTrainingStatus = async (id: string) => {
    try {
      const response = await fetch(`/api/training/status/${id}`);
      if (response.ok) {
        const data = await response.json();
        setProgress(data);
      }
    } catch (err) {
      console.error('Failed to fetch training status:', err);
    }
  };

  return {
    progress,
    error,
    isComplete,
    isConnected
  };
};