/**
 * Training React Query hooks
 * Provides data fetching, caching, and state management for training operations
 */

import React from 'react';
import { useMutation, useQuery, useQueryClient, UseQueryOptions, UseMutationOptions } from '@tanstack/react-query';
import { trainingService } from '../services/training';
import { ApiError } from '../client/apiClient';
import { useAuthStore } from '../../stores/auth.store';
import type {
  TrainingJobRequest,
  TrainingJobResponse,
  TrainingJobStatus,
  SingleProductTrainingRequest,
  ActiveModelResponse,
  ModelMetricsResponse,
  TrainedModelResponse,
  TenantStatistics,
  ModelPerformanceResponse,
  ModelsQueryParams,
  PaginatedResponse,
} from '../types/training';

// Query Keys Factory
export const trainingKeys = {
  all: ['training'] as const,
  jobs: {
    all: () => [...trainingKeys.all, 'jobs'] as const,
    status: (tenantId: string, jobId: string) => 
      [...trainingKeys.jobs.all(), 'status', tenantId, jobId] as const,
  },
  models: {
    all: () => [...trainingKeys.all, 'models'] as const,
    lists: () => [...trainingKeys.models.all(), 'list'] as const,
    list: (tenantId: string, params?: ModelsQueryParams) => 
      [...trainingKeys.models.lists(), tenantId, params] as const,
    details: () => [...trainingKeys.models.all(), 'detail'] as const,
    detail: (tenantId: string, modelId: string) => 
      [...trainingKeys.models.details(), tenantId, modelId] as const,
    active: (tenantId: string, inventoryProductId: string) => 
      [...trainingKeys.models.all(), 'active', tenantId, inventoryProductId] as const,
    metrics: (tenantId: string, modelId: string) => 
      [...trainingKeys.models.all(), 'metrics', tenantId, modelId] as const,
    performance: (tenantId: string, modelId: string) => 
      [...trainingKeys.models.all(), 'performance', tenantId, modelId] as const,
  },
  statistics: (tenantId: string) => 
    [...trainingKeys.all, 'statistics', tenantId] as const,
} as const;

// Training Job Queries
export const useTrainingJobStatus = (
  tenantId: string,
  jobId: string,
  options?: Omit<UseQueryOptions<TrainingJobStatus, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<TrainingJobStatus, ApiError>({
    queryKey: trainingKeys.jobs.status(tenantId, jobId),
    queryFn: () => trainingService.getTrainingJobStatus(tenantId, jobId),
    enabled: !!tenantId && !!jobId,
    refetchInterval: 5000, // Poll every 5 seconds while training
    staleTime: 1000, // Consider data stale after 1 second
    ...options,
  });
};

// Model Queries
export const useActiveModel = (
  tenantId: string,
  inventoryProductId: string,
  options?: Omit<UseQueryOptions<ActiveModelResponse, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<ActiveModelResponse, ApiError>({
    queryKey: trainingKeys.models.active(tenantId, inventoryProductId),
    queryFn: () => trainingService.getActiveModel(tenantId, inventoryProductId),
    enabled: !!tenantId && !!inventoryProductId,
    staleTime: 5 * 60 * 1000, // 5 minutes
    ...options,
  });
};

export const useModels = (
  tenantId: string,
  queryParams?: ModelsQueryParams,
  options?: Omit<UseQueryOptions<PaginatedResponse<TrainedModelResponse>, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<PaginatedResponse<TrainedModelResponse>, ApiError>({
    queryKey: trainingKeys.models.list(tenantId, queryParams),
    queryFn: () => trainingService.getModels(tenantId, queryParams),
    enabled: !!tenantId,
    staleTime: 2 * 60 * 1000, // 2 minutes
    ...options,
  });
};

export const useModelMetrics = (
  tenantId: string,
  modelId: string,
  options?: Omit<UseQueryOptions<ModelMetricsResponse, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<ModelMetricsResponse, ApiError>({
    queryKey: trainingKeys.models.metrics(tenantId, modelId),
    queryFn: () => trainingService.getModelMetrics(tenantId, modelId),
    enabled: !!tenantId && !!modelId,
    staleTime: 10 * 60 * 1000, // 10 minutes
    ...options,
  });
};

export const useModelPerformance = (
  tenantId: string,
  modelId: string,
  options?: Omit<UseQueryOptions<ModelPerformanceResponse, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<ModelPerformanceResponse, ApiError>({
    queryKey: trainingKeys.models.performance(tenantId, modelId),
    queryFn: () => trainingService.getModelPerformance(tenantId, modelId),
    enabled: !!tenantId && !!modelId,
    staleTime: 10 * 60 * 1000, // 10 minutes
    ...options,
  });
};

// Statistics Queries
export const useTenantTrainingStatistics = (
  tenantId: string,
  options?: Omit<UseQueryOptions<TenantStatistics, ApiError>, 'queryKey' | 'queryFn'>
) => {
  return useQuery<TenantStatistics, ApiError>({
    queryKey: trainingKeys.statistics(tenantId),
    queryFn: () => trainingService.getTenantStatistics(tenantId),
    enabled: !!tenantId,
    staleTime: 5 * 60 * 1000, // 5 minutes
    ...options,
  });
};

// Training Job Mutations
export const useCreateTrainingJob = (
  options?: UseMutationOptions<
    TrainingJobResponse,
    ApiError,
    { tenantId: string; request: TrainingJobRequest }
  >
) => {
  const queryClient = useQueryClient();

  return useMutation<
    TrainingJobResponse,
    ApiError,
    { tenantId: string; request: TrainingJobRequest }
  >({
    mutationFn: ({ tenantId, request }) => trainingService.createTrainingJob(tenantId, request),
    onSuccess: (data, { tenantId }) => {
      // Add the job status to cache
      queryClient.setQueryData(
        trainingKeys.jobs.status(tenantId, data.job_id),
        {
          job_id: data.job_id,
          status: data.status,
          progress: 0,
          message: data.message,
        }
      );
      
      // Invalidate statistics to reflect the new training job
      queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
    },
    ...options,
  });
};

export const useTrainSingleProduct = (
  options?: UseMutationOptions<
    TrainingJobResponse,
    ApiError,
    { tenantId: string; inventoryProductId: string; request: SingleProductTrainingRequest }
  >
) => {
  const queryClient = useQueryClient();

  return useMutation<
    TrainingJobResponse,
    ApiError,
    { tenantId: string; inventoryProductId: string; request: SingleProductTrainingRequest }
  >({
    mutationFn: ({ tenantId, inventoryProductId, request }) =>
      trainingService.trainSingleProduct(tenantId, inventoryProductId, request),
    onSuccess: (data, { tenantId, inventoryProductId }) => {
      // Add the job status to cache
      queryClient.setQueryData(
        trainingKeys.jobs.status(tenantId, data.job_id),
        {
          job_id: data.job_id,
          status: data.status,
          progress: 0,
          message: data.message,
        }
      );
      
      // Invalidate active model for this product
      queryClient.invalidateQueries({
        queryKey: trainingKeys.models.active(tenantId, inventoryProductId)
      });
      
      // Invalidate statistics
      queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
    },
    ...options,
  });
};

// Admin Mutations
export const useDeleteAllTenantModels = (
  options?: UseMutationOptions<{ message: string }, ApiError, { tenantId: string }>
) => {
  const queryClient = useQueryClient();

  return useMutation<{ message: string }, ApiError, { tenantId: string }>({
    mutationFn: ({ tenantId }) => trainingService.deleteAllTenantModels(tenantId),
    onSuccess: (_, { tenantId }) => {
      // Invalidate all model-related queries for this tenant
      queryClient.invalidateQueries({ queryKey: trainingKeys.models.all() });
      queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
    },
    ...options,
  });
};

// WebSocket Hook for Real-time Training Updates
export const useTrainingWebSocket = (
  tenantId: string,
  jobId: string,
  token?: string,
  options?: {
    onProgress?: (data: any) => void;
    onCompleted?: (data: any) => void;
    onError?: (error: any) => void;
    onStarted?: (data: any) => void;
    onCancelled?: (data: any) => void;
  }
) => {
  const queryClient = useQueryClient();
  const authToken = useAuthStore((state) => state.token);
  const [isConnected, setIsConnected] = React.useState(false);
  const [connectionError, setConnectionError] = React.useState<string | null>(null);

  // Memoize options to prevent unnecessary effect re-runs
  const memoizedOptions = React.useMemo(() => options, [
    options?.onProgress,
    options?.onCompleted,
    options?.onError,
    options?.onStarted,
    options?.onCancelled
  ]);

  React.useEffect(() => {
    if (!tenantId || !jobId || !memoizedOptions) {
      return;
    }

    let ws: WebSocket | null = null;
    let reconnectTimer: NodeJS.Timeout | null = null;
    let isManuallyDisconnected = false;
    let reconnectAttempts = 0;
    const maxReconnectAttempts = 3;

    const connect = () => {
      try {
        setConnectionError(null);
        const effectiveToken = token || authToken;
        console.log(`🔄 Attempting WebSocket connection (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts + 1}):`, {
          tenantId,
          jobId,
          hasToken: !!effectiveToken
        });
        
        ws = trainingService.createWebSocketConnection(tenantId, jobId, token || authToken || undefined);

        ws.onopen = () => {
          console.log('✅ Training WebSocket connected successfully');
          setIsConnected(true);
          reconnectAttempts = 0; // Reset on successful connection
          
          // Request current status on connection
          try {
            ws?.send('get_status');
            console.log('📤 Requested current training status');
          } catch (e) {
            console.warn('Failed to request status on connection:', e);
          }
          
          // Set up periodic ping to keep connection alive
          const pingInterval = setInterval(() => {
            if (ws?.readyState === WebSocket.OPEN && !isManuallyDisconnected) {
              try {
                ws?.send('ping');
                console.log('💓 Sent ping to server');
              } catch (e) {
                console.warn('Failed to send ping:', e);
                clearInterval(pingInterval);
              }
            } else {
              clearInterval(pingInterval);
            }
          }, 30000); // Ping every 30 seconds
          
          // Store interval for cleanup
          (ws as any).pingInterval = pingInterval;
        };

        ws.onmessage = (event) => {
          try {
            // Handle non-JSON messages (like pong responses)
            if (typeof event.data === 'string' && event.data === 'pong') {
              console.log('🏓 Pong received from server');
              return;
            }

            const message = JSON.parse(event.data);
            
            console.log('🔔 Training WebSocket message received:', message);

            // Handle heartbeat messages
            if (message.type === 'heartbeat') {
              console.log('💓 Heartbeat received from server');
              return; // Don't process heartbeats further
            }

            // Extract data from backend message structure
            const eventData = message.data || {};
            const progress = eventData.progress || 0;
            const currentStep = eventData.current_step || eventData.step_name || '';
            const statusMessage = eventData.message || eventData.status || '';

            // Update job status in cache with backend structure
            queryClient.setQueryData(
              trainingKeys.jobs.status(tenantId, jobId),
              (oldData: TrainingJobStatus | undefined) => ({
                ...oldData,
                job_id: jobId,
                status: message.type === 'completed' ? 'completed' : 
                        message.type === 'failed' ? 'failed' : 
                        message.type === 'started' ? 'running' : 
                        oldData?.status || 'running',
                progress: typeof progress === 'number' ? progress : oldData?.progress || 0,
                message: statusMessage || oldData?.message || '',
                current_step: currentStep || oldData?.current_step,
                estimated_time_remaining: eventData.estimated_time_remaining || oldData?.estimated_time_remaining,
              })
            );

            // Call appropriate callback based on message type (exact backend mapping)
            switch (message.type) {
              case 'started':
                memoizedOptions?.onStarted?.(message);
                break;
              case 'progress':
                memoizedOptions?.onProgress?.(message);
                break;
              case 'step_completed':
                memoizedOptions?.onProgress?.(message); // Treat step completion as progress
                break;
              case 'completed':
                console.log('✅ Training completed successfully');
                memoizedOptions?.onCompleted?.(message);
                // Invalidate models and statistics
                queryClient.invalidateQueries({ queryKey: trainingKeys.models.all() });
                queryClient.invalidateQueries({ queryKey: trainingKeys.statistics(tenantId) });
                isManuallyDisconnected = true; // Don't reconnect after completion
                break;
              case 'failed':
                console.log('❌ Training failed');
                memoizedOptions?.onError?.(message);
                isManuallyDisconnected = true; // Don't reconnect after failure
                break;
              case 'cancelled':
                console.log('🛑 Training cancelled');
                memoizedOptions?.onCancelled?.(message);
                isManuallyDisconnected = true; // Don't reconnect after cancellation
                break;
              case 'current_status':
                console.log('📊 Received current training status');
                // Treat current status as progress update if it has progress data
                if (message.data) {
                  memoizedOptions?.onProgress?.(message);
                }
                break;
              default:
                console.log(`🔍 Received unknown message type: ${message.type}`);
                break;
            }
          } catch (error) {
            console.error('Error parsing WebSocket message:', error);
            setConnectionError('Error parsing message from server');
          }
        };

        ws.onerror = (error) => {
          console.error('Training WebSocket error:', error);
          setConnectionError('WebSocket connection error');
          setIsConnected(false);
        };

        ws.onclose = (event) => {
          console.log(`❌ Training WebSocket disconnected. Code: ${event.code}, Reason: "${event.reason}"`);
          setIsConnected(false);

          // Detailed logging for different close codes
          switch (event.code) {
            case 1000:
              console.log('🔒 WebSocket closed normally');
              break;
            case 1006:
              console.log('⚠️ WebSocket closed abnormally (1006) - likely server-side issue or network problem');
              break;
            case 1001:
              console.log('🔄 WebSocket endpoint going away');
              break;
            case 1003:
              console.log('❌ WebSocket unsupported data received');
              break;
            default:
              console.log(`❓ WebSocket closed with code ${event.code}`);
          }

          // Try to reconnect if not manually disconnected and haven't exceeded max attempts
          if (!isManuallyDisconnected && event.code !== 1000 && reconnectAttempts < maxReconnectAttempts) {
            const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), 10000); // Exponential backoff, max 10s
            console.log(`🔄 Attempting to reconnect WebSocket in ${delay/1000}s... (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts})`);
            
            reconnectTimer = setTimeout(() => {
              reconnectAttempts++;
              connect();
            }, delay);
          } else if (reconnectAttempts >= maxReconnectAttempts) {
            console.log(`❌ Max reconnection attempts (${maxReconnectAttempts}) reached. Giving up.`);
            setConnectionError(`Connection failed after ${maxReconnectAttempts} attempts. The training job may not exist or the server may be unavailable.`);
          }
        };

      } catch (error) {
        console.error('Error creating WebSocket connection:', error);
        setConnectionError('Failed to create WebSocket connection');
      }
    };

    // Delay initial connection to ensure training job is created
    const initialConnectionTimer = setTimeout(() => {
      console.log('🚀 Starting initial WebSocket connection...');
      connect();
    }, 2000); // 2-second delay to let the job initialize

    // Cleanup function
    return () => {
      isManuallyDisconnected = true;
      
      if (initialConnectionTimer) {
        clearTimeout(initialConnectionTimer);
      }
      
      if (reconnectTimer) {
        clearTimeout(reconnectTimer);
      }
      
      if (ws) {
        ws.close(1000, 'Component unmounted');
      }
      
      setIsConnected(false);
    };
  }, [tenantId, jobId, token, authToken, queryClient, memoizedOptions]);

  return {
    isConnected,
    connectionError
  };
};

// Utility Hooks
export const useIsTrainingInProgress = (tenantId: string, jobId?: string) => {
  const { data: jobStatus } = useTrainingJobStatus(tenantId, jobId || '', {
    enabled: !!jobId,
  });

  return jobStatus?.status === 'running' || jobStatus?.status === 'pending';
};

export const useTrainingProgress = (tenantId: string, jobId?: string) => {
  const { data: jobStatus } = useTrainingJobStatus(tenantId, jobId || '', {
    enabled: !!jobId,
  });

  return {
    progress: jobStatus?.progress || 0,
    currentStep: jobStatus?.current_step,
    estimatedTimeRemaining: jobStatus?.estimated_time_remaining,
    isComplete: jobStatus?.status === 'completed',
    isFailed: jobStatus?.status === 'failed',
    isRunning: jobStatus?.status === 'running',
  };
};