"""
Tests for download tracking functionality.
"""

import pytest
from unittest.mock import Mock, patch

from hypersets.download_tracking import (
    DataSavings,
    DownloadStats,
    estimate_data_usage,
    create_download_stats,
)
from hypersets.dataset_info import DatasetInfo


class TestDataSavings:
    """Test DataSavings functionality."""

    def test_data_savings_calculation(self):
        """Test data savings calculation and string representation."""
        savings = DataSavings(
            total_dataset_size_gb=100.0,
            estimated_downloaded_gb=10.0,
            savings_gb=90.0,
            savings_percentage=90.0,
        )

        assert savings.total_dataset_size_gb == 100.0
        assert savings.estimated_downloaded_gb == 10.0
        assert savings.savings_gb == 90.0
        assert savings.savings_percentage == 90.0
        assert "saved you 90.0 GB (90.0% of dataset)" in str(savings)


class TestDownloadStats:
    """Test DownloadStats functionality."""

    def test_download_stats_properties(self):
        """Test DownloadStats data savings property."""
        dataset_info = DatasetInfo(
            name="test/dataset", total_parquet_files=20
        )  # 20 * 50MB = ~1GB

        stats = DownloadStats(
            query_time_seconds=2.5,
            rows_returned=1000,
            columns_returned=5,
            estimated_data_processed_mb=100.0,  # 0.1GB processed
            dataset_info=dataset_info,
        )

        savings = stats.data_savings
        assert savings.total_dataset_size_gb == pytest.approx(0.976, rel=1e-2)  # ~1GB
        assert savings.estimated_downloaded_gb == pytest.approx(
            0.098, rel=1e-2
        )  # ~0.1GB
        assert savings.savings_gb > 0
        assert savings.savings_percentage > 0


class TestDataUsageEstimation:
    """Test data usage estimation logic."""

    def test_estimate_limit_query(self):
        """Test estimation for LIMIT queries."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=50)

        estimated_mb = estimate_data_usage(
            dataset_info=dataset_info,
            query="SELECT * FROM dataset LIMIT 100",
            result_rows=100,
            result_columns=10,
            config=None,
            split=None,
        )

        # LIMIT queries should estimate fewer files accessed
        # Should be less than full dataset processing
        assert estimated_mb > 0
        assert estimated_mb < 50 * 50 * 0.5  # Less than half the dataset

    def test_estimate_sample_query_small(self):
        """Test estimation for small SAMPLE queries."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=30)

        estimated_mb = estimate_data_usage(
            dataset_info=dataset_info,
            query="SELECT * FROM dataset USING SAMPLE 500",
            result_rows=500,
            result_columns=8,
            config=None,
            split=None,
        )

        assert estimated_mb > 0
        # Small samples should access fewer files
        assert estimated_mb < 30 * 50 * 0.5  # Less than half

    def test_estimate_sample_query_large(self):
        """Test estimation for large SAMPLE queries."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=40)

        estimated_mb = estimate_data_usage(
            dataset_info=dataset_info,
            query="SELECT * FROM dataset USING SAMPLE 10000",
            result_rows=10000,
            result_columns=12,
            config=None,
            split=None,
        )

        assert estimated_mb > 0
        # Large samples might need more files

    def test_estimate_aggregation_query(self):
        """Test estimation for aggregation queries."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=20)

        estimated_mb = estimate_data_usage(
            dataset_info=dataset_info,
            query="SELECT COUNT(*) FROM dataset",
            result_rows=1,
            result_columns=1,
            config=None,
            split=None,
        )

        assert estimated_mb > 0
        # Aggregation might need full scan
        # Should be higher than simple LIMIT queries

    def test_estimate_with_config_split_filter(self):
        """Test estimation with config/split filters."""
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=100)

        # Mock the get_parquet_urls to return fewer files for filtered case
        with patch.object(
            dataset_info,
            "get_parquet_urls",
            return_value=["file1.parquet", "file2.parquet"],
        ):
            estimated_mb = estimate_data_usage(
                dataset_info=dataset_info,
                query="SELECT * FROM dataset LIMIT 1000",
                result_rows=1000,
                result_columns=5,
                config="specific_config",
                split="train",
            )

        assert estimated_mb > 0
        # Should be based on filtered file count, not total


class TestCreateDownloadStats:
    """Test create_download_stats function."""

    @patch("hypersets.download_tracking.get_dataset_info")
    def test_create_download_stats(self, mock_get_info):
        """Test creating download stats."""
        # Mock dataset info with proper structure
        mock_dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=25)
        # Add a mock config so get_parquet_urls doesn't fail
        from hypersets.dataset_info import ConfigInfo, SplitInfo

        config_info = ConfigInfo(name="en")
        split_info = SplitInfo(
            name="train", file_count=5, file_paths=["file1.parquet", "file2.parquet"]
        )
        config_info.splits["train"] = split_info
        mock_dataset_info.configs["en"] = config_info
        mock_get_info.return_value = mock_dataset_info

        stats = create_download_stats(
            dataset_name="test/dataset",
            query="SELECT title FROM dataset LIMIT 100",
            query_time=1.5,
            result_rows=100,
            result_columns=1,
            config="en",
            split="train",
            token="test_token",
        )

        # Verify mock was called correctly
        mock_get_info.assert_called_once_with("test/dataset", token="test_token")

        # Verify stats structure
        assert isinstance(stats, DownloadStats)
        assert stats.query_time_seconds == 1.5
        assert stats.rows_returned == 100
        assert stats.columns_returned == 1
        assert stats.estimated_data_processed_mb > 0
        assert stats.dataset_info is mock_dataset_info

        # Should be able to calculate savings
        savings = stats.data_savings
        assert isinstance(savings, DataSavings)
        assert savings.total_dataset_size_gb > 0

    @patch("hypersets.download_tracking.estimate_data_usage")
    @patch("hypersets.download_tracking.get_dataset_info")
    def test_create_download_stats_calls_estimate(self, mock_get_info, mock_estimate):
        """Test that create_download_stats calls estimate_data_usage correctly."""
        mock_dataset_info = DatasetInfo(name="test/dataset")
        mock_get_info.return_value = mock_dataset_info
        mock_estimate.return_value = 123.45

        stats = create_download_stats(
            dataset_name="test/dataset",
            query="SELECT * FROM dataset USING SAMPLE 1000",
            query_time=2.0,
            result_rows=1000,
            result_columns=10,
            config="config1",
            split="validation",
        )

        # Verify estimate_data_usage was called with correct parameters
        mock_estimate.assert_called_once_with(
            dataset_info=mock_dataset_info,
            query="SELECT * FROM dataset USING SAMPLE 1000",
            result_rows=1000,
            result_columns=10,
            config="config1",
            split="validation",
        )

        # Verify the estimated value was used
        assert stats.estimated_data_processed_mb == 123.45
