frontend/src/components/Evaluation/hooks/useEvaluation.js

import { useState, useRef, useEffect } from "react";
import API_CONFIG from "../../../config/api";

// Define all evaluation steps in sequence
const EVALUATION_STEPS = [
  "initializing",
  "finding_available_model_providers",
  "starting_evaluation_process",
  "evaluating_models",
  "storing_evaluation_results",
];

// Step labels for display
const STEP_LABELS = {
  initializing: "Initializing evaluation environment",
  finding_available_model_providers: "Finding available model providers",
  starting_evaluation_process: "Starting evaluation process",
  evaluating_models: "Evaluating models",
  storing_evaluation_results: "Storing evaluation results",
};

// Error messages that should be treated as errors
const ERROR_MESSAGES = [
  "heavy load",
  "try again later",
  "rate limit",
  "RATE_LIMIT_EXCEEDED",
];

export const useEvaluation = (sessionId, onComplete) => {
  const [error, setError] = useState(null);
  const [evaluationComplete, setEvaluationComplete] = useState(false);
  const [currentStep, setCurrentStep] = useState(0);
  const [evaluationStarted, setEvaluationStarted] = useState(false);
  const pollingIntervalRef = useRef(null);

  const mapStepToIndex = (step) => {
    return EVALUATION_STEPS.indexOf(step);
  };

  const checkForErrors = (logs) => {
    if (!logs) return false;

    const hasError = ERROR_MESSAGES.some((errorMessage) =>
      logs.some((log) => log.toLowerCase().includes(errorMessage.toLowerCase()))
    );

    if (hasError) {
      setError(
        "The demo is currently under heavy load, please try again later."
      );
      setEvaluationComplete(true);
      if (pollingIntervalRef.current) {
        clearInterval(pollingIntervalRef.current);
      }
      return true;
    }
    return false;
  };

  const startEvaluation = async () => {
    if (!sessionId) {
      setError("Missing session ID");
      return;
    }

    setEvaluationStarted(true);

    try {
      const response = await fetch(
        `${API_CONFIG.BASE_URL}/evaluate-benchmark`,
        {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
          },
          body: JSON.stringify({
            session_id: sessionId,
          }),
        }
      );

      const result = await response.json();

      if (response.ok) {
        setupPolling();
      } else {
        setError(result.error || "Benchmark evaluation failed");
      }
    } catch (error) {
      console.error("Error starting evaluation:", error);
      setError("Error connecting to server");
    }
  };

  const setupPolling = () => {
    pollingIntervalRef.current = setInterval(async () => {
      try {
        const logsResponse = await fetch(
          `${API_CONFIG.BASE_URL}/evaluation-logs/${sessionId}`
        );

        if (logsResponse.ok) {
          const logsResult = await logsResponse.json();

          // Check for error messages in logs
          if (checkForErrors(logsResult.logs)) {
            return;
          }

          if (logsResult.is_completed) {
            setEvaluationComplete(true);
            clearInterval(pollingIntervalRef.current);
            if (onComplete) {
              onComplete();
            }
          } else if (logsResult.current_step) {
            const newStepIndex = mapStepToIndex(logsResult.current_step);
            if (newStepIndex !== -1) {
              setCurrentStep(newStepIndex);
            }
          }
        }
      } catch (error) {
        console.log("Error polling logs:", error);
      }
    }, 2000);
  };

  useEffect(() => {
    return () => {
      if (pollingIntervalRef.current) {
        clearInterval(pollingIntervalRef.current);
      }
    };
  }, []);

  return {
    error,
    evaluationComplete,
    currentStep,
    evaluationStarted,
    startEvaluation,
    currentStepLabel:
      STEP_LABELS[EVALUATION_STEPS[currentStep]] || "Processing",
    totalSteps: EVALUATION_STEPS.length,
  };
};