src/hooks/useTraceUploader.ts

import { useCallback, useRef, useState } from 'react';
import { uploadTraceToModal } from '@/services/api';
import { getTraceExportData } from '@/services/jsonExporter';
import { AgentTrace, AgentStep, AgentTraceMetadata, FinalStep } from '@/types/agent';

/**
 * Callback function that returns fresh trace data at upload time.
 * This ensures we capture the latest state including user_evaluation.
 */
type GetTraceDataFn = () => {
  trace?: AgentTrace;
  steps: AgentStep[];
  metadata?: AgentTraceMetadata;
  finalStep?: FinalStep;
};

interface UseTraceUploaderOptions {
  getTraceData: GetTraceDataFn;
}

interface UseTraceUploaderReturn {
  uploadTrace: (forceUpload?: boolean) => Promise<boolean>;
  isUploading: boolean;
  uploadError: string | null;
  uploadSuccess: boolean;
  hasUploaded: boolean;
  resetUploadState: () => void;
}

/**
 * Custom hook to upload trace data to Modal's trace storage endpoint.
 * Uses a callback to get fresh data at upload time, ensuring user evaluations are captured.
 * 
 * Upload behavior:
 * - First upload happens automatically when task completes
 * - Subsequent uploads (with evaluation) will overwrite the previous trace
 * - Use forceUpload=true to always upload regardless of previous upload state
 */
export const useTraceUploader = ({
  getTraceData,
}: UseTraceUploaderOptions): UseTraceUploaderReturn => {
  const [isUploading, setIsUploading] = useState(false);
  const [uploadError, setUploadError] = useState<string | null>(null);
  const [uploadSuccess, setUploadSuccess] = useState(false);
  const hasUploadedRef = useRef(false);
  const uploadCountRef = useRef(0);

  const uploadTrace = useCallback(async (forceUpload: boolean = false): Promise<boolean> => {
    // Get fresh data at upload time
    const { trace, steps, metadata, finalStep } = getTraceData();
    
    // Don't upload if no trace
    if (!trace) {
      console.log('Skipping trace upload: no trace');
      return false;
    }

    // Don't upload if trace is still running
    if (trace.isRunning) {
      console.log('Skipping trace upload: trace is still running');
      return false;
    }

    // Get the export data with the latest state (including user evaluation)
    const traceData = getTraceExportData(trace, steps, metadata, finalStep);
    
    // Check if this upload has user evaluation
    const currentEvaluation = (traceData as { user_evaluation?: string }).user_evaluation;
    const hasEvaluation = currentEvaluation && currentEvaluation !== 'not_evaluated';
    
    // Always allow upload if:
    // 1. forceUpload is true (explicit request to upload with new evaluation)
    // 2. Never uploaded before
    // 3. Has new evaluation data
    if (!forceUpload && hasUploadedRef.current && !hasEvaluation) {
      console.log('Skipping trace upload: already uploaded and no new evaluation');
      return false;
    }

    setIsUploading(true);
    setUploadError(null);
    uploadCountRef.current += 1;
    const uploadNum = uploadCountRef.current;

    try {
      console.log(`Uploading trace to Modal (upload #${uploadNum})...`, { 
        traceId: trace.id, 
        user_evaluation: currentEvaluation,
        forceUpload,
        hasEvaluation 
      });
      
      const result = await uploadTraceToModal(traceData);
      
      if (result.success) {
        console.log(`✅ Trace uploaded successfully (upload #${uploadNum}):`, result);
        hasUploadedRef.current = true;
        setUploadSuccess(true);
        return true;
      } else {
        console.warn(`⚠️ Trace upload failed (upload #${uploadNum}):`, result.error);
        setUploadError(result.error || 'Unknown error');
        return false;
      }
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : 'Unknown error';
      console.error(`❌ Error uploading trace (upload #${uploadNum}):`, error);
      setUploadError(errorMessage);
      return false;
    } finally {
      setIsUploading(false);
    }
  }, [getTraceData]);

  const resetUploadState = useCallback(() => {
    hasUploadedRef.current = false;
    uploadCountRef.current = 0;
    setUploadSuccess(false);
    setUploadError(null);
  }, []);

  return {
    uploadTrace,
    isUploading,
    uploadError,
    uploadSuccess,
    hasUploaded: hasUploadedRef.current,
    resetUploadState,
  };
};