"""Tests for YDB connection module."""

import asyncio
import sys
import unittest
from unittest.mock import ANY, AsyncMock, MagicMock, patch

# Add mocks for mcp.server.handler
from tests.mocks import MockRequestHandler, mock_register_handler

# Mock the imports
sys.modules["mcp.server"] = MagicMock()
sys.modules["mcp.server.handler"] = MagicMock()
sys.modules["mcp.server.handler"].RequestHandler = MockRequestHandler
sys.modules["mcp.server.handler"].register_handler = mock_register_handler

# Import after mocking
from ydb_mcp.connection import YDBConnection  # noqa: E402


class TestYDBConnection(unittest.TestCase):
    """Test cases for YDBConnection class."""

    def test_extract_database_path(self):
        """Test database path extraction from connection string."""
        # Test with simple format
        conn = YDBConnection("grpc://ydb.server:2136/local")
        self.assertEqual(conn._extract_database_path(), "/local")

        # Test with path containing multiple segments
        conn = YDBConnection("grpc://ydb.server:2136/my/database/path")
        self.assertEqual(conn._extract_database_path(), "/my/database/path")

        # Test with query parameters
        conn = YDBConnection("grpc://ydb.server:2136/local?ssl=true&timeout=60")
        self.assertEqual(conn._extract_database_path(), "/local")

        # Test with database:// prefix
        conn = YDBConnection("database://ydb.server:2136/local")
        self.assertEqual(conn._extract_database_path(), "/local")

    @patch("ydb.aio.Driver")
    async def test_connect(self, mock_driver_class):
        """Test connection establishment."""
        # Setup mocks
        mock_driver = AsyncMock()
        mock_driver.wait = AsyncMock(return_value=True)
        mock_driver.discovery_debug_details = MagicMock(return_value="Resolved endpoints: localhost:2136")
        mock_driver_class.return_value = mock_driver

        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
            # Setup session pool mock
            mock_session_pool = MagicMock()
            mock_session_pool_class.return_value = mock_session_pool

            # Create connection and connect
            conn = YDBConnection("grpc://ydb.server:2136/local")
            driver, pool = await conn.connect()

            # Verify driver was created with correct parameters
            mock_driver_class.assert_called_once()
            mock_driver.wait.assert_called_once()
            mock_driver.discovery_debug_details.assert_called()

            # Verify session pool was created
            mock_session_pool_class.assert_called_once()

            # Verify driver and session pool were stored and returned
            assert conn.driver == mock_driver
            assert conn.session_pool == mock_session_pool
            assert driver == mock_driver
            assert pool == mock_session_pool

            # Reset mocks for next test
            mock_driver_class.reset_mock()
            mock_driver.wait.reset_mock()
            mock_driver.discovery_debug_details.reset_mock()

    @patch("ydb.aio.Driver")
    async def test_connect_with_database_in_endpoint(self, mock_driver_class):
        """Test connection with database specified in endpoint."""
        # Setup mocks
        mock_driver = AsyncMock()
        mock_driver.wait = AsyncMock(return_value=True)
        mock_driver.discovery_debug_details = MagicMock(return_value="Resolved endpoints: localhost:2136")
        mock_driver_class.return_value = mock_driver

        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
            # Setup session pool mock
            mock_session_pool = MagicMock()
            mock_session_pool_class.return_value = mock_session_pool

            # Test cases for different endpoint formats
            test_cases = [
                ("grpc://ydb.server:2136/local", "grpc://ydb.server:2136", "/local"),
                ("grpcs://ydb.server:2136/local/test", "grpcs://ydb.server:2136", "/local/test"),
                ("ydb.server:2136/local", "grpc://ydb.server:2136", "/local"),
                ("grpc://ydb.server:2136/local", "grpc://ydb.server:2136", "/local"),
            ]

            for endpoint, expected_endpoint, expected_database in test_cases:
                # Create connection and connect
                conn = YDBConnection(endpoint)
                await conn.connect()

                # Verify driver was created with correct parameters
                mock_driver_class.assert_called_with(
                    endpoint=expected_endpoint, database=expected_database, credentials=ANY
                )
                mock_driver.wait.assert_called_once()
                mock_driver.discovery_debug_details.assert_called()

                # Reset mock call count
                mock_driver_class.reset_mock()
                mock_driver.wait.reset_mock()
                mock_driver.discovery_debug_details.reset_mock()

    @patch("ydb.aio.Driver")
    async def test_connect_with_explicit_database(self, mock_driver_class):
        """Test connection with explicitly provided database."""
        # Setup mocks
        mock_driver = AsyncMock()
        mock_driver.wait = AsyncMock(return_value=True)
        mock_driver.discovery_debug_details = MagicMock(return_value="Resolved endpoints: localhost:2136")
        mock_driver_class.return_value = mock_driver

        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
            # Setup session pool mock
            mock_session_pool = MagicMock()
            mock_session_pool_class.return_value = mock_session_pool

            # Test cases for different endpoint formats with explicit database
            test_cases = [
                (
                    "grpc://ydb.server:2136/local",
                    "/explicit",
                    "grpc://ydb.server:2136",
                    "/explicit",
                ),
                (
                    "grpcs://ydb.server:2136/local",
                    "explicit",
                    "grpcs://ydb.server:2136",
                    "/explicit",
                ),
                ("ydb.server:2136/local", "/other", "grpc://ydb.server:2136", "/other"),
            ]

            for endpoint, database, expected_endpoint, expected_database in test_cases:
                # Create connection and connect
                conn = YDBConnection(endpoint, database=database)
                await conn.connect()

                # Verify driver was created with correct parameters
                mock_driver_class.assert_called_with(
                    endpoint=expected_endpoint, database=expected_database, credentials=ANY
                )
                mock_driver.wait.assert_called_once()
                mock_driver.discovery_debug_details.assert_called()

                # Reset mock call count
                mock_driver_class.reset_mock()
                mock_driver.wait.reset_mock()
                mock_driver.discovery_debug_details.reset_mock()


# Allow tests to run with asyncio
def run_async_test(test_case, test_func):
    """Run an async test function."""
    loop = asyncio.get_event_loop()
    loop.run_until_complete(test_func(test_case))


# Patch test methods to run with asyncio
for method_name in dir(TestYDBConnection):
    if method_name.startswith("test_") and method_name != "test_extract_database_path":
        method = getattr(TestYDBConnection, method_name)
        if asyncio.iscoroutinefunction(method):
            setattr(TestYDBConnection, method_name, lambda self, m=method: run_async_test(self, m))

if __name__ == "__main__":
    unittest.main()
