frontend/src/components/BenchmarkGenerator.jsx

import React, { useState, useEffect, useRef } from "react";
import { Box, Typography, CircularProgress, Alert, Paper } from "@mui/material";
import PlayArrowIcon from "@mui/icons-material/PlayArrow";
import AccessTimeIcon from "@mui/icons-material/AccessTime";
import LogDisplay from "./LogDisplay";
import { useNavigate, useSearchParams } from "react-router-dom";
import API_CONFIG from "../config/api";

// Temps de simulation en millisecondes pour les documents précalculés
const SIMULATION_DURATION = 80000; // 20 secondes

// Define all benchmark steps in sequence
const BENCHMARK_STEPS = [
  "configuration",
  "provider_check",
  "ingestion",
  "upload_ingest_to_hub",
  "summarization",
  "chunking",
  "single_shot_question_generation",
];

// Step labels for display (more user-friendly names)
const STEP_LABELS = {
  configuration: "Configuration",
  provider_check: "Finding providers",
  ingestion: "Ingestion",
  upload_ingest_to_hub: "Upload to Hub",
  summarization: "Summarization",
  chunking: "Chunking",
  single_shot_question_generation: "Question generation",
  evaluation_provider_check: "Checking evaluation providers",
  evaluation: "Running evaluations",
  evaluation_saving_results: "Saving evaluation results",
};

// Simulated log messages for pre-calculated documents
const SIMULATED_LOGS = [
  "[INFO] Initializing benchmark generation...",
  "[INFO] Generating base configuration file...",
  "[SUCCESS] Stage completed: configuration",
  "[INFO] Finding available providers for models...",
  "[SUCCESS] Stage completed: provider_check",
  "[INFO] Starting ingestion process...",
  "[SUCCESS] Stage completed: ingestion",
  "[INFO] Processing document content for upload...",
  "[SUCCESS] Stage completed: upload_ingest_to_hub",
  "[INFO] Generating document summary...",
  "[SUCCESS] Stage completed: summarization",
  "[INFO] Chunking content for better analysis...",
  "[SUCCESS] Stage completed: chunking",
  "[INFO] Generating single-shot questions...",
  "[SUCCESS] Stage completed: single_shot_question_generation",
  "[SUCCESS] Benchmark process completed successfully",
];

/**
 * Component to handle benchmark generation and display logs
 *
 * @param {Object} props - Component props
 * @param {string} props.sessionId - The session ID for the uploaded file
 * @param {boolean} props.isDefaultDocument - Whether this is a pre-calculated document
 * @param {Function} props.onComplete - Function to call when generation is complete
 * @returns {JSX.Element} Benchmark generator component
 */
const BenchmarkGenerator = ({ sessionId, isDefaultDocument, onComplete }) => {
  const [searchParams] = useSearchParams();
  const isDefault =
    searchParams.get("isDefault") === "true" || isDefaultDocument;
  const [generating, setGenerating] = useState(false);
  const [generationComplete, setGenerationComplete] = useState(false);
  const [generationLogs, setGenerationLogs] = useState([]);
  const [error, setError] = useState(null);
  const [currentPhase, setCurrentPhase] = useState("initializing");
  const [completedSteps, setCompletedSteps] = useState([]);
  const [activeStep, setActiveStep] = useState(1);
  const [elapsedTime, setElapsedTime] = useState(0);

  // Reference to keep track of the polling interval
  const pollingIntervalRef = useRef(null);

  // Reference to keep track of the timer interval
  const timerIntervalRef = useRef(null);

  // Reference for starting time
  const startTimeRef = useRef(null);

  // Simulation interval reference
  const simulationIntervalRef = useRef(null);

  // Start generation on component mount
  useEffect(() => {
    // Set start time
    startTimeRef.current = Date.now();

    // Référence pour le timeout
    let timeoutRef = null;

    // Start timer
    timerIntervalRef.current = setInterval(() => {
      const timeElapsed = Math.floor(
        (Date.now() - startTimeRef.current) / 1000
      );
      setElapsedTime(timeElapsed);

      // Vérifier si le temps écoulé dépasse 8 minutes (480 secondes) et que nous ne sommes pas en mode simulation
      if (timeElapsed > 480 && !isDefault && !generationComplete) {
        // Afficher un message d'erreur en cas de timeout
        setError(
          "The benchmark generation is taking too long. The demo is currently under heavy load, please try again later."
        );
        setGenerationComplete(true);

        // Nettoyer les intervalles
        if (pollingIntervalRef.current) {
          clearInterval(pollingIntervalRef.current);
        }

        if (timerIntervalRef.current) {
          clearInterval(timerIntervalRef.current);
        }
      }
    }, 1000);

    // Gestionnaire pour détecter quand la page redevient visible
    const handleVisibilityChange = () => {
      if (
        document.visibilityState === "visible" &&
        !isDefault &&
        !generationComplete
      ) {
        console.log("Page became visible, checking for missed steps...");
        // Force une nouvelle requête pour récupérer les logs
        const checkCurrentState = async () => {
          try {
            // D'abord essayer de récupérer les logs de benchmark
            const logsResponse = await fetch(
              `${API_CONFIG.BASE_URL}/benchmark-logs/${sessionId}`
            );

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

              // Si la tâche est terminée, mettre à jour l'état
              if (logsResult.is_completed) {
                setGenerationComplete(true);
                if (pollingIntervalRef.current) {
                  clearInterval(pollingIntervalRef.current);
                }
                if (onComplete) {
                  onComplete({
                    success: true,
                    sessionId,
                    logs: logsResult.logs,
                  });
                }
              }
            } else {
              // Si la tâche de benchmark n'existe pas, essayer les logs de configuration
              const configResponse = await fetch(
                `${API_CONFIG.BASE_URL}/config-logs/${sessionId}`
              );

              if (configResponse.ok) {
                const configResult = await configResponse.json();
                if (configResult.logs) {
                  setGenerationLogs(configResult.logs);
                }
              }
            }
          } catch (error) {
            console.error("Error checking for missed steps:", error);
          }
        };

        checkCurrentState();
      }
    };

    // Ajouter l'écouteur pour le changement de visibilité
    document.addEventListener("visibilitychange", handleVisibilityChange);

    if (isDefault) {
      simulateGeneration();
    } else {
      generateBenchmark();
    }

    // Clean up the polling interval and timer when the component unmounts
    return () => {
      if (pollingIntervalRef.current) {
        clearInterval(pollingIntervalRef.current);
      }
      if (timerIntervalRef.current) {
        clearInterval(timerIntervalRef.current);
      }
      if (simulationIntervalRef.current) {
        clearInterval(simulationIntervalRef.current);
      }
      document.removeEventListener("visibilitychange", handleVisibilityChange);
    };
  }, [isDefault, sessionId, generationComplete, onComplete]);

  // Simulate the benchmark generation for pre-calculated documents
  const simulateGeneration = () => {
    setGenerating(true);
    setGenerationLogs([]);
    setError(null);
    setCurrentPhase("initializing");
    setCompletedSteps([]);
    setActiveStep(1);

    // Timing variables for simulation
    const totalSteps = SIMULATED_LOGS.length;
    const totalDuration = SIMULATION_DURATION; // 20 seconds
    const intervalPerStep = totalDuration / totalSteps;
    let currentStep = 0;

    // Function to add next log message
    const addNextLog = () => {
      if (currentStep < SIMULATED_LOGS.length) {
        const newLogs = [...generationLogs, SIMULATED_LOGS[currentStep]];
        setGenerationLogs(newLogs);
        currentStep++;

        // Check if completed
        if (currentStep >= SIMULATED_LOGS.length) {
          // Simulation complete
          setTimeout(() => {
            setCurrentPhase("complete");
            setGenerationComplete(true);
            clearInterval(simulationIntervalRef.current);
            if (onComplete) {
              onComplete({
                success: true,
                sessionId,
                logs: newLogs,
              });
            }
          }, 1000);
        }
      }
    };

    // Start simulation
    simulationIntervalRef.current = setInterval(addNextLog, intervalPerStep);
  };

  // Determine the current phase and completed steps based on logs
  useEffect(() => {
    if (generationLogs.length === 0) return;

    // Recalculer complètement les étapes complétées à chaque fois
    // au lieu de simplement ajouter les nouvelles étapes
    const newCompletedSteps = [];

    // Identifier toutes les étapes complétées dans tous les logs
    generationLogs.forEach((log) => {
      const match = log.match(/\[SUCCESS\] Stage completed: (\w+)/);
      if (match && match[1]) {
        const completedStep = match[1].trim();
        if (
          BENCHMARK_STEPS.includes(completedStep) &&
          !newCompletedSteps.includes(completedStep)
        ) {
          newCompletedSteps.push(completedStep);
        }
      }
    });

    // Déterminer l'étape active basée sur les étapes complétées
    let newActiveStep = activeStep;

    if (newCompletedSteps.length > 0) {
      // Trouver l'étape la plus avancée dans les logs
      const maxCompletedStepIndex = Math.max(
        ...newCompletedSteps.map((step) => BENCHMARK_STEPS.indexOf(step))
      );
      // Passer à l'étape suivante
      const calculatedStep = maxCompletedStepIndex + 1;

      // Ne mettre à jour que si la nouvelle étape est plus avancée que l'étape actuelle
      if (calculatedStep > activeStep) {
        newActiveStep = calculatedStep;
      }

      // S'assurer que l'activeStep ne dépasse pas le nombre total d'étapes
      if (newActiveStep >= BENCHMARK_STEPS.length) {
        newActiveStep = BENCHMARK_STEPS.length;
      }
    } else if (activeStep === 0) {
      // Si aucune étape n'est trouvée et l'étape active est 0, passer à 1
      newActiveStep = 1;
    }

    // Mettre à jour l'état si les étapes ont changé
    if (JSON.stringify(newCompletedSteps) !== JSON.stringify(completedSteps)) {
      setCompletedSteps(newCompletedSteps);
    }

    // Mettre à jour l'étape active seulement si elle a changé
    if (newActiveStep !== activeStep) {
      setActiveStep(newActiveStep);
    }

    // Skip the rest of the log processing if we're simulating
    if (isDefault) return;

    // Check the latest logs to determine the current phase
    const recentLogs = generationLogs.slice(-10); // Check more logs

    // Detect completion conditions
    const isComplete =
      recentLogs.some((log) =>
        log.includes("[SUCCESS] Benchmark process completed successfully")
      ) ||
      recentLogs.some((log) =>
        log.includes("[SUCCESS] Ingestion process completed successfully")
      ) ||
      newCompletedSteps.includes("single_shot_question_generation") ||
      newActiveStep >= BENCHMARK_STEPS.length;

    if (isComplete) {
      setCurrentPhase("complete");
      setGenerationComplete(true);
      // Stop polling when benchmark is complete
      if (pollingIntervalRef.current) {
        clearInterval(pollingIntervalRef.current);
      }
      // Notify parent component that generation is complete
      if (onComplete) {
        console.log("Notifying parent that generation is complete");
        onComplete({
          success: true,
          sessionId,
          logs: generationLogs,
        });
      }
    } else if (
      recentLogs.some((log) => log.includes("Starting ingestion process"))
    ) {
      setCurrentPhase("benchmarking");
    } else if (
      recentLogs.some((log) => log.includes("Generating base configuration"))
    ) {
      setCurrentPhase("configuring");
    }
  }, [
    generationLogs,
    completedSteps,
    activeStep,
    sessionId,
    onComplete,
    isDefault,
  ]);

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

    setGenerating(true);
    setGenerationLogs([]);
    setError(null);
    setCurrentPhase("initializing");
    setCompletedSteps([]);
    setActiveStep(1);

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

      const result = await response.json();

      if (response.ok) {
        setGenerationLogs(result.logs || []);

        // Configurer le polling pour suivre la progression
        pollingIntervalRef.current = setInterval(async () => {
          // Vérifier si on a déjà terminé
          if (generationComplete) {
            clearInterval(pollingIntervalRef.current);
            return;
          }

          try {
            // Appeler l'API pour obtenir les derniers logs
            const logsResponse = await fetch(
              `${API_CONFIG.BASE_URL}/benchmark-progress/${sessionId}`
            );

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

              // Mettre à jour les logs s'il y en a de nouveaux
              if (
                logsResult.logs &&
                logsResult.logs.length > generationLogs.length
              ) {
                setGenerationLogs(logsResult.logs);
              }

              // Vérifier si la tâche est terminée
              if (logsResult.is_completed) {
                setGenerationComplete(true);
                clearInterval(pollingIntervalRef.current);
                // La notification est maintenant gérée dans le useEffect ci-dessus
              }
            }
          } catch (error) {
            console.log("Error polling for logs:", error);
            // Ne pas arrêter le polling en cas d'erreurs réseau
          }
        }, 2000); // Interroger toutes les 2 secondes
      } else {
        // Handle error
        setGenerationLogs([`Error: ${result.error || "Unknown error"}`]);
        setError(result.error || "Benchmark generation failed");
      }
    } catch (error) {
      console.error("Error generating benchmark:", error);
      setGenerationLogs([`Error: ${error.message || "Unknown error"}`]);
      setError("Server connection error");
    } finally {
      setGenerating(false);
    }
  };

  // Get title based on current phase
  const getPhaseTitle = () => {
    switch (currentPhase) {
      case "initializing":
        return "Benchmark generation...";
      case "configuring":
        return "Creating benchmark...";
      case "benchmarking":
        return "Creating benchmark...";
      case "complete":
        return "Benchmark generated successfully!";
      default:
        return "Processing...";
    }
  };

  // Get the current step information for display
  const getCurrentStepInfo = () => {
    const totalSteps = BENCHMARK_STEPS.length;
    const currentStepIndex = activeStep;

    // If there's no active step yet
    if (currentStepIndex <= 1 && completedSteps.length === 0) {
      return `Starting (1/${totalSteps})`;
    }

    // If all steps are completed
    if (currentStepIndex >= totalSteps) {
      return `Complete (${totalSteps}/${totalSteps})`;
    }

    // Get current step name
    const currentStepName =
      STEP_LABELS[BENCHMARK_STEPS[currentStepIndex]] || "Processing";

    return `${currentStepName} (${currentStepIndex}/${totalSteps})`;
  };

  // Format elapsed time in HH:MM:SS
  const formatElapsedTime = () => {
    const hours = Math.floor(elapsedTime / 3600);
    const minutes = Math.floor((elapsedTime % 3600) / 60);
    const seconds = elapsedTime % 60;

    return [
      hours.toString().padStart(2, "0"),
      minutes.toString().padStart(2, "0"),
      seconds.toString().padStart(2, "0"),
    ].join(":");
  };

  // If complete, stop the timer
  useEffect(() => {
    if (generationComplete && timerIntervalRef.current) {
      clearInterval(timerIntervalRef.current);
    }
  }, [generationComplete]);

  return (
    <Paper
      elevation={3}
      sx={{
        p: 4,
        mt: 3,
        mb: 3,
        display: "flex",
        flexDirection: "column",
        alignItems: "center",
        justifyContent: "center",
        minHeight: 200,
        position: "relative",
      }}
    >
      {/* Temps estimé */}
      <Box
        sx={{
          position: "absolute",
          top: 12,
          right: 12,
          backgroundColor: "rgba(0, 0, 0, 0.04)",
          borderRadius: "4px",
          px: 1,
          py: 0.5,
          display: "inline-flex",
          alignItems: "center",
        }}
      >
        <Typography
          variant="caption"
          sx={{
            fontSize: "0.675rem",
            color: "text.secondary",
            fontWeight: 500,
          }}
        >
          Estimated time: ~2 min
        </Typography>
      </Box>

      {error ? (
        <Alert severity="error" sx={{ width: "100%" }}>
          {error}
        </Alert>
      ) : (
        <>
          <CircularProgress size={60} sx={{ mb: 2 }} />
          <Typography variant="h6" component="div" gutterBottom>
            {getPhaseTitle()}
          </Typography>

          {/* Step progress indicator */}
          <Typography variant="body1" color="text.secondary">
            {getCurrentStepInfo()}
          </Typography>

          {/* Timer display */}
          <Box
            sx={{
              display: "flex",
              alignItems: "center",
              mt: 1,
              color: "text.secondary",
            }}
          >
            <Typography variant="body2" sx={{ opacity: 0.5 }}>
              {formatElapsedTime()}
            </Typography>
          </Box>
        </>
      )}

      {/* Use the LogDisplay component */}
      {/* <LogDisplay logs={generationLogs} height={150} /> */}
    </Paper>
  );
};

export default BenchmarkGenerator;