import os
import tempfile
import unittest
from unittest.mock import MagicMock, patch

from datasets import Dataset

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.multistep.config import AgentRunResult
from screensuite.benchmarks.multistep.gaia.benchmark import GaiaBenchmark
from screensuite.benchmarks.multistep.gaia.config import GaiaConfig


class TestGaiaBenchmark(unittest.TestCase):
    """Test cases for the GaiaBenchmark class."""

    def setUp(self):
        """Set up test fixtures."""
        # Create a mock config
        self.config = GaiaConfig(
            hf_repo="GeekAgents/GAIA-web",
            split="validation",
        )

        # Create a benchmark instance
        self.benchmark = GaiaBenchmark(
            name="test-benchmark",
            config=self.config,
            tags=["test"],
        )

        # Create a mock dataset
        self.mock_dataset = Dataset.from_list(
            [
                {
                    "Question": "What is the capital of France?",
                    "Final answer": "Paris",
                },
                {
                    "Question": "What is the capital of Germany?",
                    "Final answer": "Berlin",
                },
            ]
        )

    @patch("screensuite.benchmarks.hubbasebenchmark.HubBaseBenchmark.load")
    def test_load(self, mock_load_dataset):
        """Test the load method."""

        # Set up the mock to assign self.mock_dataset to self.dataset when called
        def side_effect(*args, **kwargs) -> None:
            self.benchmark.dataset = self.mock_dataset

        mock_load_dataset.side_effect = side_effect

        # Call the method
        self.benchmark.load()

        # Check that the result is the mock dataset
        assert isinstance(self.benchmark.dataset, Dataset)
        assert self.benchmark.dataset.column_names == ["question", "reference_answer"]
        assert self.benchmark.dataset["question"][0] == self.mock_dataset["Question"][0]
        assert self.benchmark.dataset["reference_answer"][0] == self.mock_dataset["Final answer"][0]

    @patch("screensuite.benchmarks.hubbasebenchmark.HubBaseBenchmark.load")
    @patch("screensuite.benchmarks.multistep.generation.answer_single_question")
    def test_evaluate(self, mock_answer_single_question, mock_load):
        """Test the evaluate method."""
        # Create a mock model
        mock_model = MagicMock()

        # Create a temporary directory for testing
        with tempfile.TemporaryDirectory() as temp_dir:
            # Create output directory structure
            output_dir = os.path.join(temp_dir, "output", "gaia")
            os.makedirs(output_dir, exist_ok=True)

            # Change working directory to temp directory
            original_dir = os.getcwd()
            os.chdir(temp_dir)

            try:
                benchmark = GaiaBenchmark(
                    name="test-benchmark",
                    config=self.config,
                    tags=["test"],
                )

                mock_load.return_value = None
                benchmark.dataset = self.mock_dataset

                # Create mock responses
                mock_response = AgentRunResult(
                    model_id="test-model",
                    question="What is the capital of France?",
                    original_question="What is the capital of France?",
                    answer="Paris",
                    reference_answer="Paris",
                    intermediate_steps=[],
                    start_time=0.0,
                    end_time=1.0,
                    token_counts={"total": 100},
                )
                mock_answer_single_question.return_value = mock_response

                # Call the evaluate method with a run_name
                run_name = "test_run"
                evaluation_config = EvaluationConfig(parallel_workers=2, test_mode=False, run_name=run_name)
                result = benchmark.evaluate(mock_model, evaluation_config)

                # Check that _get_responses was called with the correct arguments
                self.assertEqual(mock_answer_single_question.call_count, 2)
                mock_answer_single_question.assert_any_call(
                    {"Question": "What is the capital of France?", "Final answer": "Paris"},
                    mock_model,
                    f"./output/gaia/{run_name}",
                    True,
                    False,
                )
                mock_answer_single_question.assert_any_call(
                    {"Question": "What is the capital of Germany?", "Final answer": "Berlin"},
                    mock_model,
                    f"./output/gaia/{run_name}",
                    True,
                    False,
                )

                # Check the evaluation results
                self.assertEqual(result["avg_accuracy"], 1.0)  # One correct, one wrong answer

                # Check that output files were created
                expected_output_dir = os.path.join(output_dir, run_name)
                assert os.path.exists(expected_output_dir), f"Output directory {expected_output_dir} was not created"

                expected_output_file = os.path.join(expected_output_dir, "answers.jsonl")
                assert os.path.exists(expected_output_file), f"Output file {expected_output_file} was not created"

                # Check that the output file contains the expected number of entries
                with open(expected_output_file, "r") as f:
                    lines = f.readlines()
                    assert len(lines) == 2

            finally:
                # Restore original working directory
                os.chdir(original_dir)


if __name__ == "__main__":
    unittest.main()
