"""
Tests for dataset info retrieval functionality.
"""

import pytest
import tempfile
import json
from pathlib import Path
from unittest.mock import Mock, patch

from hypersets.dataset_info import (
    DatasetInfo,
    ConfigInfo,
    SplitInfo,
    DatasetInfoCache,
    _expand_glob_pattern,
    _parse_readme_yaml,
    discover_dataset_info,
    get_dataset_info,
    list_configs,
    clear_cache,
)
from hypersets.exceptions import (
    ConfigNotFoundError,
    DatasetNotFoundError,
    SplitNotFoundError,
)


class TestDatasetInfoCache:
    """Test dataset info caching."""

    def test_cache_miss(self):
        """Test cache miss returns None."""
        with tempfile.TemporaryDirectory() as temp_dir:
            cache = DatasetInfoCache(cache_dir=temp_dir, ttl_hours=24)
            result = cache.get("nonexistent/dataset")
            assert result is None

    def test_cache_set_and_get(self):
        """Test setting and getting from cache."""
        with tempfile.TemporaryDirectory() as temp_dir:
            cache = DatasetInfoCache(cache_dir=temp_dir, ttl_hours=24)

            # Create test dataset info
            dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=5)
            config_info = ConfigInfo(name="default")
            split_info = SplitInfo(
                name="train",
                file_count=3,
                file_paths=[
                    "train-00001.parquet",
                    "train-00002.parquet",
                    "train-00003.parquet",
                ],
            )
            config_info.splits["train"] = split_info
            dataset_info.configs["default"] = config_info

            # Cache and retrieve
            cache.set(dataset_info)
            retrieved = cache.get("test/dataset")

            assert retrieved is not None
            assert retrieved.name == "test/dataset"
            assert retrieved.total_parquet_files == 5
            assert "default" in retrieved.configs
            assert "train" in retrieved.configs["default"].splits
            assert retrieved.configs["default"].splits["train"].file_count == 3

    def test_cache_expiry(self):
        """Test cache expiry."""
        with tempfile.TemporaryDirectory() as temp_dir:
            cache = DatasetInfoCache(
                cache_dir=temp_dir, ttl_hours=0
            )  # Immediate expiry

            dataset_info = DatasetInfo(name="test/dataset")
            cache.set(dataset_info)

            # Should be expired immediately
            result = cache.get("test/dataset")
            assert result is None

    def test_cache_clear(self):
        """Test cache clearing."""
        with tempfile.TemporaryDirectory() as temp_dir:
            cache = DatasetInfoCache(cache_dir=temp_dir, ttl_hours=24)

            dataset_info = DatasetInfo(name="test/dataset")
            cache.set(dataset_info)

            # Verify cached
            assert cache.get("test/dataset") is not None

            # Clear and verify
            cache.clear()
            assert cache.get("test/dataset") is None


class TestDatasetInfoStructures:
    """Test dataset info data structures."""

    def test_dataset_info_properties(self):
        """Test DatasetInfo properties."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=10)

        config1 = ConfigInfo(name="config1")
        config2 = ConfigInfo(name="config2")
        dataset_info.configs["config1"] = config1
        dataset_info.configs["config2"] = config2

        assert dataset_info.config_names == ["config1", "config2"]
        assert (
            dataset_info.estimated_total_size_gb == (10 * 50) / 1024
        )  # 10 files * 50MB

    def test_config_info_properties(self):
        """Test ConfigInfo properties."""
        config = ConfigInfo(name="test_config")

        split1 = SplitInfo(
            name="train", file_count=5, file_paths=["f1", "f2", "f3", "f4", "f5"]
        )
        split2 = SplitInfo(name="test", file_count=2, file_paths=["f6", "f7"])

        config.splits["train"] = split1
        config.splits["test"] = split2

        assert config.split_names == ["train", "test"]
        assert config.total_files == 7

    def test_get_parquet_urls(self):
        """Test getting parquet URLs."""
        dataset_info = DatasetInfo(name="test/dataset")
        config = ConfigInfo(name="default")
        split = SplitInfo(
            name="train",
            file_count=2,
            file_paths=["data/train-00001.parquet", "data/train-00002.parquet"],
        )
        config.splits["train"] = split
        dataset_info.configs["default"] = config

        # Test all files
        urls = dataset_info.get_parquet_urls()
        assert len(urls) == 2
        assert all(
            "train-00001.parquet" in url or "train-00002.parquet" in url for url in urls
        )

        # Test specific config/split
        urls = dataset_info.get_parquet_urls(config="default", split="train")
        assert len(urls) == 2

        # Test invalid config
        with pytest.raises(ConfigNotFoundError, match="Config 'invalid' not found"):
            dataset_info.get_parquet_urls(config="invalid")

        # Test invalid split
        with pytest.raises(SplitNotFoundError, match="Split 'invalid' not found"):
            dataset_info.get_parquet_urls(config="default", split="invalid")


class TestDatasetInfoDiscovery:
    """Test dataset info discovery with mocked API calls."""

    @patch("hypersets.dataset_info._parse_readme_yaml")
    @patch("hypersets.dataset_info.HfApi")
    def test_discover_dataset_info_basic(self, mock_hf_api, mock_parse_yaml):
        """Test basic dataset info discovery."""
        # Mock API response
        mock_api_instance = Mock()
        mock_api_instance.list_repo_files.return_value = [
            "config1/train-00001.parquet",
            "config1/train-00002.parquet",
            "config1/test-00001.parquet",
            "config2/train-00001.parquet",
            "README.md",
            "other_file.txt",
        ]
        mock_hf_api.return_value = mock_api_instance

        # Mock YAML frontmatter
        mock_parse_yaml.return_value = {
            "configs": [
                {
                    "config_name": "config1",
                    "data_files": [
                        {"split": "train", "path": "config1/train-*"},
                        {"split": "test", "path": "config1/test-*"},
                    ],
                },
                {
                    "config_name": "config2",
                    "data_files": [{"split": "train", "path": "config2/train-*"}],
                },
            ]
        }

        # Test discovery
        dataset_info = get_dataset_info("test/dataset", use_cache=False)

        assert dataset_info.name == "test/dataset"
        assert dataset_info.total_parquet_files == 4
        assert len(dataset_info.configs) == 2
        assert "config1" in dataset_info.configs
        assert "config2" in dataset_info.configs

        # Check config1 structure
        config1 = dataset_info.configs["config1"]
        assert len(config1.splits) == 2
        assert "train" in config1.splits
        assert "test" in config1.splits
        assert config1.splits["train"].file_count == 2
        assert config1.splits["test"].file_count == 1

        # Check config2 structure
        config2 = dataset_info.configs["config2"]
        assert len(config2.splits) == 1
        assert "train" in config2.splits
        assert config2.splits["train"].file_count == 1

    @patch("hypersets.dataset_info.HfApi")
    def test_discover_no_parquet_files(self, mock_hf_api):
        """Test discovery with no parquet files."""
        mock_api_instance = Mock()
        mock_api_instance.list_repo_files.return_value = [
            "README.md",
            "config.json",
            "data.csv",
        ]
        mock_hf_api.return_value = mock_api_instance

        with pytest.raises(DatasetNotFoundError, match="No parquet files found"):
            get_dataset_info("test/dataset", use_cache=False)

    @patch("hypersets.dataset_info._parse_readme_yaml")
    @patch("hypersets.dataset_info.HfApi")
    def test_list_configs(self, mock_hf_api, mock_parse_yaml):
        """Test listing configs."""
        mock_api_instance = Mock()
        mock_api_instance.list_repo_files.return_value = [
            "en/train.parquet",
            "fr/train.parquet",
            "de/test.parquet",
            "README.md",
        ]
        mock_hf_api.return_value = mock_api_instance

        # Mock YAML frontmatter
        mock_parse_yaml.return_value = {
            "configs": [
                {
                    "config_name": "en",
                    "data_files": [{"split": "train", "path": "en/train.parquet"}],
                },
                {
                    "config_name": "fr",
                    "data_files": [{"split": "train", "path": "fr/train.parquet"}],
                },
                {
                    "config_name": "de",
                    "data_files": [{"split": "test", "path": "de/test.parquet"}],
                },
            ]
        }

        configs = list_configs("test/dataset", use_cache=False)
        assert sorted(configs) == ["de", "en", "fr"]

    def test_clear_cache_function(self):
        """Test the clear_cache function."""
        # This should not raise an error
        clear_cache()
        assert True  # If we get here, the function worked
