"""Tests for the experiment module."""

import asyncio
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
from pydantic import BaseModel

from ragas.backends.inmemory import InMemoryBackend
from ragas.experiment import Experiment, experiment, version_experiment
from ragas.experimental.dataset import Dataset
from ragas.utils import find_git_root, memorable_names


# Test data models
class SampleDataRow(BaseModel):
    question: str
    answer: str
    score: float


class ExperimentResultRow(BaseModel):
    question: str
    processed_answer: str
    sentiment: str
    processing_time: float


# Test fixtures
@pytest.fixture
def temp_dir():
    """Create a temporary directory for testing."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        yield Path(tmp_dir)


@pytest.fixture
def mock_git_repo(temp_dir):
    """Create a mock git repository."""
    git_dir = temp_dir / ".git"
    git_dir.mkdir()

    # Mock git.Repo
    mock_repo = MagicMock()
    mock_repo.is_dirty.return_value = False
    mock_repo.head.commit.hexsha = "abc123def456"
    mock_repo.git.add = MagicMock()
    mock_repo.index.commit = MagicMock()
    mock_repo.create_head = MagicMock()

    with patch("git.Repo", return_value=mock_repo):
        yield mock_repo, temp_dir


@pytest.fixture
def sample_dataset():
    """Create a sample dataset for testing."""
    backend = InMemoryBackend()
    dataset = Dataset(
        name="test_dataset",
        data_model=SampleDataRow,
        backend=backend,
        data=[
            SampleDataRow(
                question="What is Python?", answer="A programming language", score=0.9
            ),
            SampleDataRow(
                question="What is AI?", answer="Artificial Intelligence", score=0.8
            ),
            SampleDataRow(
                question="What is ML?", answer="Machine Learning", score=0.85
            ),
        ],
    )
    return dataset


@pytest.fixture
def experiment_backend():
    """Create a backend for experiments."""
    return InMemoryBackend()


# Test classes
class TestExperiment:
    """Test the Experiment class."""

    def test_experiment_inheritance(self):
        """Test that Experiment properly inherits from DataTable."""
        assert hasattr(Experiment, "DATATABLE_TYPE")
        assert Experiment.DATATABLE_TYPE == "Experiment"

    def test_experiment_creation(self, experiment_backend):
        """Test creating an Experiment instance."""
        experiment = Experiment(
            name="test_experiment",
            data_model=ExperimentResultRow,
            backend=experiment_backend,
        )

        assert experiment.name == "test_experiment"
        assert experiment.backend == experiment_backend
        assert len(experiment) == 0


class TestVersionExperiment:
    """Test the version_experiment function."""

    def test_version_experiment_no_changes(self, mock_git_repo):
        """Test version_experiment when there are no changes."""
        mock_repo, temp_dir = mock_git_repo

        # Mock that repo is clean
        mock_repo.is_dirty.return_value = False

        with patch("ragas.utils.find_git_root", return_value=temp_dir):
            commit_hash = version_experiment("test_experiment")

        assert commit_hash == "abc123def456"
        mock_repo.is_dirty.assert_called()
        mock_repo.create_head.assert_called_with(
            "ragas/test_experiment", "abc123def456"
        )

    def test_version_experiment_with_changes(self, mock_git_repo):
        """Test version_experiment when there are changes to commit."""
        mock_repo, temp_dir = mock_git_repo

        # Mock that repo is dirty
        mock_repo.is_dirty.return_value = True

        # Mock commit object
        mock_commit = MagicMock()
        mock_commit.hexsha = "new123commit456"
        mock_repo.index.commit.return_value = mock_commit

        with patch("ragas.utils.find_git_root", return_value=temp_dir):
            commit_hash = version_experiment("test_experiment")

        assert commit_hash == "new123commit456"
        mock_repo.git.add.assert_called_with("-u")
        mock_repo.index.commit.assert_called_once()

    def test_version_experiment_with_custom_message(self, mock_git_repo):
        """Test version_experiment with custom commit message."""
        mock_repo, temp_dir = mock_git_repo
        mock_repo.is_dirty.return_value = True

        mock_commit = MagicMock()
        mock_commit.hexsha = "custom123commit456"
        mock_repo.index.commit.return_value = mock_commit

        with patch("ragas.utils.find_git_root", return_value=temp_dir):
            version_experiment(
                "test_experiment", commit_message="Custom experiment message"
            )

        mock_repo.index.commit.assert_called_with("Custom experiment message")

    def test_version_experiment_stage_all(self, mock_git_repo):
        """Test version_experiment with stage_all=True."""
        mock_repo, temp_dir = mock_git_repo
        mock_repo.is_dirty.return_value = True

        mock_commit = MagicMock()
        mock_commit.hexsha = "staged123commit456"
        mock_repo.index.commit.return_value = mock_commit

        with patch("ragas.utils.find_git_root", return_value=temp_dir):
            version_experiment("test_experiment", stage_all=True)

        mock_repo.git.add.assert_called_with(".")

    def test_version_experiment_no_branch_creation(self, mock_git_repo):
        """Test version_experiment with create_branch=False."""
        mock_repo, temp_dir = mock_git_repo

        with patch("ragas.utils.find_git_root", return_value=temp_dir):
            version_experiment("test_experiment", create_branch=False)

        mock_repo.create_head.assert_not_called()

    def test_find_git_root_error_handling(self, temp_dir):
        """Test that find_git_root raises ValueError when no git repo found."""
        with pytest.raises(ValueError, match="No git repository found"):
            find_git_root(temp_dir)


class TestExperimentDecorator:
    """Test the experiment decorator."""

    @pytest.mark.asyncio
    async def test_simple_async_experiment(self, sample_dataset, experiment_backend):
        """Test a simple async experiment function."""

        @experiment(experiment_model=ExperimentResultRow, backend=experiment_backend)
        async def simple_experiment(row: SampleDataRow) -> ExperimentResultRow:
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer.upper(),
                sentiment="positive",
                processing_time=0.1,
            )

        # Test that decorator creates proper wrapper
        assert hasattr(simple_experiment, "arun")
        assert hasattr(simple_experiment, "__call__")

        # Test calling the wrapped function directly
        test_row = SampleDataRow(question="Test?", answer="test answer", score=0.5)
        result = await simple_experiment(test_row)

        assert isinstance(result, ExperimentResultRow)
        assert result.processed_answer == "TEST ANSWER"
        assert result.sentiment == "positive"

    @pytest.mark.asyncio
    async def test_experiment_arun(self, sample_dataset, experiment_backend):
        """Test running experiment against a dataset."""

        @experiment(experiment_model=ExperimentResultRow, backend=experiment_backend)
        async def test_experiment(row: SampleDataRow) -> ExperimentResultRow:
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer.lower(),
                sentiment="neutral",
                processing_time=0.05,
            )

        # Mock memorable_names to return predictable name
        with patch(
            "ragas.utils.memorable_names.generate_unique_name",
            return_value="test_experiment_name",
        ):
            experiment_result = await test_experiment.arun(sample_dataset)

        assert isinstance(experiment_result, Experiment)
        assert experiment_result.name == "test_experiment_name"
        assert len(experiment_result) == 3  # Should have processed all 3 items

    @pytest.mark.asyncio
    async def test_experiment_with_name_prefix(
        self, sample_dataset, experiment_backend
    ):
        """Test experiment decorator with name prefix."""

        @experiment(
            experiment_model=ExperimentResultRow,
            backend=experiment_backend,
            name_prefix="prefix",
        )
        async def prefixed_experiment(row: SampleDataRow) -> ExperimentResultRow:
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer,
                sentiment="neutral",
                processing_time=0.01,
            )

        with patch(
            "ragas.utils.memorable_names.generate_unique_name",
            return_value="random_name",
        ):
            experiment_result = await prefixed_experiment.arun(sample_dataset)

        assert experiment_result.name == "prefix-random_name"

    @pytest.mark.asyncio
    async def test_experiment_with_custom_name(
        self, sample_dataset, experiment_backend
    ):
        """Test experiment with custom name."""

        @experiment(experiment_model=ExperimentResultRow, backend=experiment_backend)
        async def custom_named_experiment(row: SampleDataRow) -> ExperimentResultRow:
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer,
                sentiment="positive",
                processing_time=0.02,
            )

        experiment_result = await custom_named_experiment.arun(
            sample_dataset, name="my_custom_experiment"
        )

        assert experiment_result.name == "my_custom_experiment"

    def test_sync_experiment_function(self, experiment_backend):
        """Test that sync functions work with the experiment decorator."""

        @experiment(experiment_model=ExperimentResultRow, backend=experiment_backend)
        def sync_experiment(row: SampleDataRow) -> ExperimentResultRow:
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer.upper(),
                sentiment="positive",
                processing_time=0.0,
            )

        # Test that we can call it synchronously within async context
        test_row = SampleDataRow(question="Sync test?", answer="sync answer", score=0.7)

        async def test_sync_call():
            result = await sync_experiment(test_row)
            return result

        result = asyncio.run(test_sync_call())
        assert isinstance(result, ExperimentResultRow)
        assert result.processed_answer == "SYNC ANSWER"

    @pytest.mark.asyncio
    async def test_experiment_error_handling(self, sample_dataset, experiment_backend):
        """Test that experiment handles individual task failures gracefully."""

        @experiment(experiment_model=ExperimentResultRow, backend=experiment_backend)
        async def failing_experiment(row: SampleDataRow) -> ExperimentResultRow:
            if "AI" in row.question:  # Fail on the AI question
                raise ValueError("Test error")
            return ExperimentResultRow(
                question=row.question,
                processed_answer=row.answer,
                sentiment="neutral",
                processing_time=0.01,
            )

        # Should continue processing other items even if some fail
        with patch(
            "ragas.utils.memorable_names.generate_unique_name",
            return_value="error_test",
        ):
            experiment_result = await failing_experiment.arun(sample_dataset)

        # Should have 2 successful results (3 items - 1 failure)
        assert len(experiment_result) == 2

    @pytest.mark.asyncio
    async def test_experiment_with_no_model(self, sample_dataset, experiment_backend):
        """Test experiment without specifying a model."""

        @experiment(backend=experiment_backend)
        async def untyped_experiment(row: SampleDataRow) -> dict:
            return {"question": row.question, "answer": row.answer, "processed": True}

        with patch(
            "ragas.utils.memorable_names.generate_unique_name",
            return_value="untyped_test",
        ):
            experiment_result = await untyped_experiment.arun(sample_dataset)

        assert isinstance(experiment_result, Experiment)
        assert len(experiment_result) == 3


class TestMemorableNames:
    """Test the memorable names functionality."""

    def test_memorable_names_generation(self):
        """Test that memorable names are generated correctly."""
        name = memorable_names.generate_name()
        assert "_" in name
        parts = name.split("_", 1)  # Split on first underscore only
        assert len(parts) == 2
        assert parts[0] in memorable_names.adjectives
        assert parts[1] in memorable_names.scientists

    def test_unique_name_generation(self):
        """Test that unique names are generated."""
        # Create a fresh instance to avoid state from other tests
        from ragas.utils import MemorableNames

        generator = MemorableNames()

        names = [generator.generate_unique_name() for _ in range(10)]
        assert len(set(names)) == 10  # All names should be unique

    def test_unique_names_batch_generation(self):
        """Test batch generation of unique names."""
        from ragas.utils import MemorableNames

        generator = MemorableNames()

        names = generator.generate_unique_names(5)
        assert len(names) == 5
        assert len(set(names)) == 5  # All should be unique


class TestUtilityFunctions:
    """Test utility functions added to ragas.utils."""

    def test_find_git_root_with_git_repo(self, temp_dir):
        """Test find_git_root finds git repository correctly."""
        # Create a nested directory structure with .git at the top
        git_dir = temp_dir / ".git"
        git_dir.mkdir()

        nested_dir = temp_dir / "nested" / "deeply" / "nested"
        nested_dir.mkdir(parents=True)

        # Should find git root from nested directory
        found_root = find_git_root(nested_dir)
        # Use resolve() to handle symlinks and get canonical path
        assert found_root.resolve() == temp_dir.resolve()

    def test_find_git_root_current_dir(self):
        """Test find_git_root uses current directory when no path provided."""
        # This should find the actual git root of the ragas project
        try:
            root = find_git_root()
            assert isinstance(root, Path)
            assert (root / ".git").exists()
        except ValueError:
            # If we're not in a git repo, that's expected
            pass

    def test_find_git_root_no_repo_error(self, temp_dir):
        """Test find_git_root raises error when no git repo found."""
        with pytest.raises(ValueError, match="No git repository found"):
            find_git_root(temp_dir)
