"""
Tests for memory configuration functionality.
"""

import pytest
from unittest.mock import Mock, patch
import duckdb

from hypersets.duckdb_mount import DuckDBMount
from hypersets.query_interface import query


class TestMemoryConfiguration:
    """Test memory configuration in DuckDB connections."""

    def test_default_memory_limit(self):
        """Test that default memory limit is 4GB."""
        mount = DuckDBMount()
        assert mount.memory_limit == "4GB"

    def test_custom_memory_limit(self):
        """Test setting custom memory limit."""
        mount = DuckDBMount(memory_limit="8GB")
        assert mount.memory_limit == "8GB"

    @patch("hypersets.duckdb_mount.duckdb.connect")
    def test_memory_limit_applied_to_connection(self, mock_connect):
        """Test that memory limit is actually applied to DuckDB connection."""
        mock_conn = Mock()
        mock_connect.return_value = mock_conn

        mount = DuckDBMount(memory_limit="2GB")
        mount.connect()

        # Verify memory_limit was set
        expected_calls = [
            "INSTALL httpfs",
            "LOAD httpfs",
            "SET memory_limit='2GB'",
            "SET threads=4",
            "SET http_timeout=300000",
            "SET http_retries=3",
            "SET http_retry_wait_ms=5000",
            "SET enable_progress_bar=false",
        ]

        actual_calls = [call[0][0] for call in mock_conn.execute.call_args_list]
        for expected_call in expected_calls:
            assert expected_call in actual_calls

    def test_threads_configuration(self):
        """Test threads configuration."""
        mount = DuckDBMount(threads=8)
        assert mount.threads == 8

    @patch("hypersets.duckdb_mount.duckdb.connect")
    def test_memory_limit_with_real_connection(self, mock_connect):
        """Test memory limit setting with a real DuckDB connection mock."""
        # Create a real mock connection that behaves like DuckDB
        mock_conn = Mock()
        mock_connect.return_value = mock_conn

        mount = DuckDBMount(memory_limit="6GB", threads=2)
        mount.connect()

        # Verify the connection was created correctly
        mock_connect.assert_called_once_with(":memory:")

        # Check that memory_limit was set
        execute_calls = [call[0][0] for call in mock_conn.execute.call_args_list]
        assert "SET memory_limit='6GB'" in execute_calls
        assert "SET threads=2" in execute_calls

    @patch("hypersets.query_interface.mount_dataset")
    def test_query_passes_memory_limit(self, mock_mount):
        """Test that query function passes memory_limit to mount_dataset."""
        mock_mount_instance = Mock()
        mock_mount_instance.query.return_value = Mock()
        mock_mount.__enter__ = Mock(return_value=mock_mount_instance)
        mock_mount.__exit__ = Mock(return_value=None)

        # Mock info
        with patch("hypersets.query_interface.info") as mock_get_info:
            mock_get_info.return_value = Mock()

            # Call query with custom memory limit
            query(
                "SELECT * FROM dataset LIMIT 10",
                dataset="test/dataset",
                memory_limit="8GB",
                threads=6,
            )

            # Verify mount_dataset was called with correct parameters
            mock_mount.assert_called_once()
            call_kwargs = mock_mount.call_args[1]
            assert call_kwargs["memory_limit"] == "8GB"
            assert call_kwargs["threads"] == 6

    def test_memory_limit_validation(self):
        """Test that memory limit accepts valid formats."""
        # Test various valid formats
        valid_limits = ["1GB", "2048MB", "4GB", "8192MB", "512MB"]

        for limit in valid_limits:
            mount = DuckDBMount(memory_limit=limit)
            assert mount.memory_limit == limit

    def test_context_manager_memory_config(self):
        """Test memory configuration with context manager."""
        from hypersets.duckdb_mount import mount_dataset

        # Mock the actual DuckDB connection since we're just testing parameter passing
        with patch("hypersets.duckdb_mount.DuckDBMount") as mock_mount_class:
            mock_mount_instance = Mock()
            mock_mount_class.return_value = mock_mount_instance

            # Use context manager with memory config
            with mount_dataset(
                "test/dataset", memory_limit="16GB", threads=12
            ) as mount:
                pass

            # Verify DuckDBMount was created with correct parameters
            mock_mount_class.assert_called_once_with(memory_limit="16GB", threads=12)


class TestMemoryConfigurationIntegration:
    """Integration tests for memory configuration with mocked datasets."""

    @patch("hypersets.query_interface.info")
    @patch("hypersets.duckdb_mount.DuckDBMount")
    def test_memory_limit_end_to_end(self, mock_mount_class, mock_get_info):
        """Test memory limit configuration end-to-end."""
        # Mock dataset info
        mock_get_info.return_value = Mock()

        # Mock DuckDB mount
        mock_mount_instance = Mock()
        mock_mount_instance.query.return_value = Mock(
            columns=["col1"], values=[["val1"]]
        )
        mock_mount_class.return_value = mock_mount_instance

        # Test with custom memory limit
        result = query(
            "SELECT * FROM dataset LIMIT 1", dataset="test/dataset", memory_limit="12GB"
        )

        # Verify mount was created with correct memory limit
        mock_mount_class.assert_called_once()
        init_kwargs = mock_mount_class.call_args[1]
        assert init_kwargs["memory_limit"] == "12GB"

        # Verify the result was created
        assert result is not None
