"""
Tests for query interface functionality.
"""

import pytest
import pandas as pd
from unittest.mock import Mock, patch, MagicMock

from hypersets.query_interface import QueryResult, query, sample, head, count, schema
from hypersets.dataset_info import DatasetInfo, ConfigInfo, SplitInfo


class TestQueryResult:
    """Test QueryResult functionality."""

    def test_basic_query_result(self):
        """Test basic QueryResult operations."""
        df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
        dataset_info = DatasetInfo(name="test/dataset")

        result = QueryResult(
            _df=df,
            query="SELECT * FROM dataset",
            dataset_name="test/dataset",
            config=None,
            split=None,
            query_time_seconds=1.5,
            dataset_info=dataset_info,
        )

        # Test basic properties
        assert len(result) == 3
        assert result.shape == (3, 2)
        assert result.columns == ["a", "b"]
        assert "rows=3" in str(result)
        assert "time=1.50s" in str(result)

        # Test to_pandas
        df_copy = result.to_pandas()
        pd.testing.assert_frame_equal(df, df_copy)

        # Test head and sample
        head_df = result.head(2)
        assert len(head_df) == 2

        sample_df = result.sample(2)
        assert len(sample_df) == 2

    def test_query_result_with_download_tracking(self):
        """Test QueryResult with download tracking enabled."""
        df = pd.DataFrame({"title": ["Article 1", "Article 2"]})
        dataset_info = DatasetInfo(name="test/dataset", total_parquet_files=10)

        # Mock the download tracking import to avoid circular imports during test
        with patch(
            "hypersets.download_tracking.create_download_stats"
        ) as mock_create_stats:
            mock_stats = Mock()
            mock_stats.data_savings = Mock()
            mock_stats.data_savings.__str__ = Mock(return_value="Saved 5.0 GB!")
            mock_create_stats.return_value = mock_stats

            result = QueryResult(
                _df=df,
                query="SELECT title FROM dataset LIMIT 2",
                dataset_name="test/dataset",
                config=None,
                split=None,
                query_time_seconds=2.0,
                dataset_info=dataset_info,
                track_downloads=True,
            )

            # Should have called create_download_stats
            mock_create_stats.assert_called_once()

            # Should have download stats
            assert hasattr(result, "_download_stats")
            assert result._download_stats is not None

            # String representation should include savings
            result_str = str(result)
            assert "Saved 5.0 GB!" in result_str

    def test_to_dataset_without_datasets_library(self):
        """Test to_dataset when datasets library not available."""
        df = pd.DataFrame({"a": [1, 2, 3]})
        dataset_info = DatasetInfo(name="test/dataset")

        result = QueryResult(
            _df=df,
            query="SELECT * FROM dataset",
            dataset_name="test/dataset",
            config=None,
            split=None,
            query_time_seconds=1.0,
            dataset_info=dataset_info,
        )

        # Mock HF_DATASETS_AVAILABLE as False
        with patch("hypersets.query_interface.HF_DATASETS_AVAILABLE", False):
            with pytest.raises(RuntimeError, match="datasets library not installed"):
                result.to_dataset()

    def test_data_savings_without_tracking(self):
        """Test accessing data_savings when tracking is disabled."""
        df = pd.DataFrame({"a": [1, 2, 3]})
        dataset_info = DatasetInfo(name="test/dataset")

        result = QueryResult(
            _df=df,
            query="SELECT * FROM dataset",
            dataset_name="test/dataset",
            config=None,
            split=None,
            query_time_seconds=1.0,
            dataset_info=dataset_info,
            track_downloads=False,
        )

        with pytest.raises(RuntimeError, match="Download tracking not enabled"):
            _ = result.data_savings


class TestQueryFunctions:
    """Test main query functions."""

    @patch("hypersets.query_interface.mount_dataset")
    @patch("hypersets.query_interface.get_dataset_info")
    def test_query_function(self, mock_get_info, mock_mount):
        """Test the main query function."""
        # Mock dataset info
        mock_dataset_info = DatasetInfo(name="test/dataset")
        mock_get_info.return_value = mock_dataset_info

        # Mock mount context manager
        mock_mount_instance = Mock()
        mock_mount_instance.query.return_value = pd.DataFrame({"col1": [1, 2, 3]})
        mock_mount.return_value.__enter__.return_value = mock_mount_instance
        mock_mount.return_value.__exit__.return_value = None

        # Execute query
        result = query(
            sql="SELECT * FROM dataset LIMIT 3",
            dataset="test/dataset",
            config="config1",
            split="train",
        )

        # Verify result
        assert isinstance(result, QueryResult)
        assert len(result) == 3
        assert result.query == "SELECT * FROM dataset LIMIT 3"
        assert result.dataset_name == "test/dataset"
        assert result.config == "config1"
        assert result.split == "train"

        # Verify mocks were called correctly
        mock_get_info.assert_called_once_with("test/dataset", token=None)
        mock_mount.assert_called_once()
        mock_mount_instance.query.assert_called_once_with(
            "SELECT * FROM dataset LIMIT 3"
        )

    @patch("hypersets.query_interface.query")
    def test_sample_function(self, mock_query):
        """Test the sample function."""
        mock_result = Mock(spec=QueryResult)
        mock_query.return_value = mock_result

        # Test basic sample
        result = sample(n=100, dataset="test/dataset")

        mock_query.assert_called_once_with(
            "SELECT * FROM dataset USING SAMPLE 100",
            dataset="test/dataset",
            config=None,
            split=None,
            track_downloads=False,
        )
        assert result is mock_result

        # Test sample with columns
        mock_query.reset_mock()
        sample(n=50, dataset="test/dataset", columns=["title", "text"])

        mock_query.assert_called_with(
            "SELECT title, text FROM dataset USING SAMPLE 50",
            dataset="test/dataset",
            config=None,
            split=None,
            track_downloads=False,
        )

    @patch("hypersets.query_interface.query")
    def test_head_function(self, mock_query):
        """Test the head function."""
        mock_result = Mock(spec=QueryResult)
        mock_query.return_value = mock_result

        result = head(n=10, dataset="test/dataset", config="config1")

        mock_query.assert_called_once_with(
            "SELECT * FROM dataset LIMIT 10",
            dataset="test/dataset",
            config="config1",
            split=None,
            track_downloads=False,
        )
        assert result is mock_result

    @patch("hypersets.query_interface.query")
    def test_count_function(self, mock_query):
        """Test the count function."""
        # Mock QueryResult with count data
        mock_df = pd.DataFrame({"count": [12345]})
        mock_result = Mock(spec=QueryResult)
        mock_result.to_pandas.return_value = mock_df
        mock_query.return_value = mock_result

        result = count(dataset="test/dataset", config="config1", split="train")

        mock_query.assert_called_once_with(
            "SELECT COUNT(*) as count FROM dataset",
            dataset="test/dataset",
            config="config1",
            split="train",
        )
        assert result == 12345

    @patch("hypersets.query_interface.query")
    def test_schema_function(self, mock_query):
        """Test the schema function."""
        mock_df = pd.DataFrame(
            {"column_name": ["col1", "col2"], "column_type": ["INTEGER", "VARCHAR"]}
        )
        mock_result = Mock(spec=QueryResult)
        mock_result.to_pandas.return_value = mock_df
        mock_query.return_value = mock_result

        result = schema(dataset="test/dataset")

        mock_query.assert_called_once_with(
            "DESCRIBE dataset", dataset="test/dataset", config=None, split=None
        )
        pd.testing.assert_frame_equal(result, mock_df)
