import os
from unittest.mock import patch

import pytest
from datasets import Dataset, DatasetDict

from screensuite.basebenchmark import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import (
    HubBaseBenchmark,
    HubBaseBenchmarkConfig,
)


class TestHubBaseBenchmarkConfig(HubBaseBenchmarkConfig):
    """Test configuration class for HubBaseBenchmark"""

    pass


class TestHubBaseBenchmark(HubBaseBenchmark[TestHubBaseBenchmarkConfig, None]):
    """Test implementation of HubBaseBenchmark"""

    def evaluate(self, model, evaluation_config, env_config: None = None) -> BenchmarkResult:
        return BenchmarkResult(
            metrics={"test": 1.0},
            reference_field="test",
        )


@pytest.fixture
def mock_dataset():
    """Create a mock dataset"""
    return Dataset.from_dict({"test": [1, 2, 3]})


@pytest.fixture
def mock_dataset_dict():
    """Create a mock dataset dictionary"""
    return DatasetDict(
        {"train": Dataset.from_dict({"test": [1, 2, 3]}), "test": Dataset.from_dict({"test": [4, 5, 6]})}
    )


@pytest.fixture
def mock_load_dataset(mock_dataset, mock_dataset_dict):
    """Mock the load_dataset function"""
    with patch("screensuite.benchmarks.hubbasebenchmark.load_dataset") as mock:

        def side_effect(*args, **kwargs):
            if isinstance(kwargs.get("split"), list):
                return mock_dataset_dict
            return mock_dataset

        mock.side_effect = side_effect
        yield mock


@pytest.fixture
def mock_login():
    """Mock the huggingface_hub login function"""
    with patch("screensuite.benchmarks.hubbasebenchmark.login") as mock:
        yield mock


def test_single_data_dir_and_split(mock_load_dataset, mock_login):
    """Test loading with single data_dir and split"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split="train", data_dir="data")
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    benchmark.load()

    assert benchmark.dataset is not None
    assert not benchmark.datasets
    mock_load_dataset.assert_called_once_with(
        "test/repo", split="train", revision="main", data_dir="data", streaming=False
    )


def test_multiple_splits(mock_load_dataset, mock_login):
    """Test loading with single data_dir and multiple splits"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split=["train", "test"], data_dir="data")
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    benchmark.load()

    assert benchmark.dataset is None
    assert benchmark.datasets
    assert len(benchmark.datasets) == 2
    assert "train" in benchmark.datasets
    assert "test" in benchmark.datasets


def test_multiple_data_dirs(mock_load_dataset, mock_login):
    """Test loading with multiple data_dirs and single split"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split="train", data_dir=["data1", "data2"])
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    benchmark.load()

    assert benchmark.dataset is None
    assert len(benchmark.datasets) == 2
    assert "data1" in benchmark.datasets
    assert "data2" in benchmark.datasets


def test_multiple_data_dirs_and_splits(mock_load_dataset, mock_login):
    """Test loading with multiple data_dirs and splits"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split=["train", "test"], data_dir=["data1", "data2"])
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    benchmark.load()

    assert benchmark.dataset is None
    assert len(benchmark.datasets) == 4
    assert "train_data1" in benchmark.datasets
    assert "test_data1" in benchmark.datasets
    assert "train_data2" in benchmark.datasets
    assert "test_data2" in benchmark.datasets


def test_hf_token_login(mock_load_dataset, mock_login):
    """Test that HF token login is called when token is present"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split="train", data_dir="data")
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    with patch.dict(os.environ, {"HF_TOKEN": "test_token"}):
        benchmark.load()

    mock_login.assert_called_once_with(token="test_token")


def test_no_hf_token(mock_load_dataset, mock_login):
    """Test that HF token login is not called when token is not present"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split="train", data_dir="data")
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    with patch.dict(os.environ, {}, clear=True):
        benchmark.load()

    mock_login.assert_not_called()


def test_reload_does_nothing(mock_load_dataset, mock_login):
    """Test that reloading does not call load_dataset again"""
    config = TestHubBaseBenchmarkConfig(hf_repo="test/repo", split="train", data_dir="data")
    benchmark = TestHubBaseBenchmark("test", config, ["test"])

    # First load
    benchmark.load()
    mock_load_dataset.reset_mock()

    # Second load
    benchmark.load()
    mock_load_dataset.assert_not_called()
