"""Tests for per-backend HTTP/2 connection pooling in BackendDispatcher.

Verifies that each backend URL gets its own ``httpx.AsyncClient`` with
configurable pool limits, that pools are reused for the same URL, and
that shutdown closes all pools.
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

from llmhosts.config import ConnectionPoolConfig
from llmhosts.proxy.dispatcher import BackendDispatcher, BackendEntry

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_backend(
    name: str = "generic-local",
    base_url: str = "http://127.0.0.1:11434",
    backend_type: str = "generic",
    models: list[str] | None = None,
) -> BackendEntry:
    """Create a BackendEntry for testing."""
    return BackendEntry(
        name=name,
        base_url=base_url,
        backend_type=backend_type,
        models=models or ["llama3.2"],
        priority=0,
        healthy=True,
    )


# ---------------------------------------------------------------------------
# Pool creation and reuse
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_created_on_first_use() -> None:
    """_get_or_create_pool creates a new pool on the first call for a backend URL."""
    config = ConnectionPoolConfig(http2=False)  # h2 may not be installed
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()

    assert len(dispatcher._pool_registry) == 0
    pool = dispatcher._get_or_create_pool(backend)

    assert pool is not None
    assert isinstance(pool, httpx.AsyncClient)
    assert backend.base_url in dispatcher._pool_registry
    assert len(dispatcher._pool_registry) == 1

    await dispatcher.shutdown()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_reused_for_same_url() -> None:
    """Same base_url returns the same AsyncClient instance."""
    config = ConnectionPoolConfig(http2=False)
    dispatcher = BackendDispatcher(pool_config=config)
    backend_a = _make_backend(name="a", base_url="http://10.0.0.1:11434")
    backend_b = _make_backend(name="b", base_url="http://10.0.0.1:11434")

    pool_a = dispatcher._get_or_create_pool(backend_a)
    pool_b = dispatcher._get_or_create_pool(backend_b)

    assert pool_a is pool_b
    assert len(dispatcher._pool_registry) == 1

    await dispatcher.shutdown()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_different_urls_get_different_pools() -> None:
    """Different backend URLs get separate connection pools."""
    config = ConnectionPoolConfig(http2=False)
    dispatcher = BackendDispatcher(pool_config=config)
    backend_a = _make_backend(name="a", base_url="http://10.0.0.1:11434")
    backend_b = _make_backend(name="b", base_url="http://10.0.0.2:11434")

    pool_a = dispatcher._get_or_create_pool(backend_a)
    pool_b = dispatcher._get_or_create_pool(backend_b)

    assert pool_a is not pool_b
    assert len(dispatcher._pool_registry) == 2

    await dispatcher.shutdown()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_registry_key_is_base_url() -> None:
    """Pool registry uses the backend base_url as the dictionary key."""
    config = ConnectionPoolConfig(http2=False)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend(base_url="http://gpu-box:11434")
    dispatcher._get_or_create_pool(backend)

    assert "http://gpu-box:11434" in dispatcher._pool_registry

    await dispatcher.shutdown()


# ---------------------------------------------------------------------------
# HTTP/2 and pool limits configuration
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.asyncio
async def test_http2_graceful_fallback_when_h2_missing() -> None:
    """When http2=True but h2 is not installed, pool falls back to HTTP/1.1."""
    config = ConnectionPoolConfig(http2=True)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()

    # This should NOT raise -- it should fall back to HTTP/1.1
    pool = dispatcher._get_or_create_pool(backend)
    assert pool is not None
    assert isinstance(pool, httpx.AsyncClient)
    assert backend.base_url in dispatcher._pool_registry

    await dispatcher.shutdown()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_http2_flag_set_on_pool() -> None:
    """Pool is created with http2=True when config says so."""
    config = ConnectionPoolConfig(http2=True)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()

    with patch("llmhosts.proxy.dispatcher.httpx.AsyncClient") as mock_cls:
        mock_cls.return_value = MagicMock()
        dispatcher._get_or_create_pool(backend)

        mock_cls.assert_called_once()
        call_kwargs = mock_cls.call_args
        assert call_kwargs.kwargs.get("http2") is True or call_kwargs[1].get("http2") is True

    # Clean up mock pool
    dispatcher._pool_registry.clear()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_http2_disabled_when_config_false() -> None:
    """Pool is created with http2=False when config disables it."""
    config = ConnectionPoolConfig(http2=False)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()

    with patch("llmhosts.proxy.dispatcher.httpx.AsyncClient") as mock_cls:
        mock_cls.return_value = MagicMock()
        dispatcher._get_or_create_pool(backend)

        call_kwargs = mock_cls.call_args
        assert call_kwargs.kwargs.get("http2") is False or call_kwargs[1].get("http2") is False

    dispatcher._pool_registry.clear()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_limits_match_config() -> None:
    """Pool limits (max_connections, max_keepalive, keepalive_expiry) come from config."""
    config = ConnectionPoolConfig(max_connections=30, max_keepalive=10, idle_timeout=60)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()

    with patch("llmhosts.proxy.dispatcher.httpx.AsyncClient") as mock_cls:
        mock_cls.return_value = MagicMock()
        dispatcher._get_or_create_pool(backend)

        call_kwargs = mock_cls.call_args
        limits = call_kwargs.kwargs.get("limits") or call_kwargs[1].get("limits")
        assert limits is not None
        assert limits.max_connections == 30
        assert limits.max_keepalive_connections == 10
        assert limits.keepalive_expiry == 60

    dispatcher._pool_registry.clear()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_default_config_values() -> None:
    """Default ConnectionPoolConfig values match the spec: 20, 5, 30, True."""
    config = ConnectionPoolConfig()
    assert config.max_connections == 20
    assert config.max_keepalive == 5
    assert config.idle_timeout == 30
    assert config.http2 is True


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_timeout_uses_backend_timeout() -> None:
    """Pool timeout is set from BACKEND_TIMEOUT constant (120s total, 10s connect)."""
    dispatcher = BackendDispatcher()
    backend = _make_backend()

    with patch("llmhosts.proxy.dispatcher.httpx.AsyncClient") as mock_cls:
        mock_cls.return_value = MagicMock()
        dispatcher._get_or_create_pool(backend)

        call_kwargs = mock_cls.call_args
        timeout = call_kwargs.kwargs.get("timeout") or call_kwargs[1].get("timeout")
        assert timeout is not None
        assert timeout.read == 120.0
        assert timeout.connect == 10.0

    dispatcher._pool_registry.clear()


# ---------------------------------------------------------------------------
# Shutdown
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.asyncio
async def test_shutdown_closes_all_pools() -> None:
    """shutdown() calls aclose() on every pool and clears the registry."""
    dispatcher = BackendDispatcher()

    mock_pool_a = AsyncMock(spec=httpx.AsyncClient)
    mock_pool_b = AsyncMock(spec=httpx.AsyncClient)

    dispatcher._pool_registry["http://a:11434"] = mock_pool_a
    dispatcher._pool_registry["http://b:11434"] = mock_pool_b

    await dispatcher.shutdown()

    mock_pool_a.aclose.assert_awaited_once()
    mock_pool_b.aclose.assert_awaited_once()
    assert len(dispatcher._pool_registry) == 0


@pytest.mark.unit
@pytest.mark.asyncio
async def test_shutdown_with_empty_registry() -> None:
    """shutdown() works gracefully when no pools were ever created."""
    dispatcher = BackendDispatcher()
    assert len(dispatcher._pool_registry) == 0
    await dispatcher.shutdown()  # should not raise
    assert len(dispatcher._pool_registry) == 0


@pytest.mark.unit
@pytest.mark.asyncio
async def test_double_shutdown_is_safe() -> None:
    """Calling shutdown() twice does not raise."""
    config = ConnectionPoolConfig(http2=False)
    dispatcher = BackendDispatcher(pool_config=config)
    backend = _make_backend()
    dispatcher._get_or_create_pool(backend)

    await dispatcher.shutdown()
    await dispatcher.shutdown()  # second call should be a no-op

    assert len(dispatcher._pool_registry) == 0


# ---------------------------------------------------------------------------
# Integration with dispatcher lifecycle
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.asyncio
async def test_startup_registers_core_when_loaded() -> None:
    """When Core engine has a model loaded, startup registers it as a backend."""
    dispatcher = BackendDispatcher()
    mock_engine = type("MockEngine", (), {"is_loaded": True})()
    dispatcher._core_engine = mock_engine

    await dispatcher.startup()

    core_backends = [b for b in dispatcher._backends if b.name == "llmhosts-core"]
    assert len(core_backends) == 1
    assert core_backends[0].backend_type == "llmhosts-core"
    assert core_backends[0].priority == -1

    await dispatcher.shutdown()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_config_passed_through_constructor() -> None:
    """ConnectionPoolConfig passed to BackendDispatcher is stored and used."""
    config = ConnectionPoolConfig(max_connections=42, http2=False)
    dispatcher = BackendDispatcher(pool_config=config)

    assert dispatcher._pool_config.max_connections == 42
    assert dispatcher._pool_config.http2 is False


@pytest.mark.unit
@pytest.mark.asyncio
async def test_default_pool_config_when_none() -> None:
    """When pool_config is None, defaults are used."""
    dispatcher = BackendDispatcher(pool_config=None)

    assert dispatcher._pool_config.max_connections == 20
    assert dispatcher._pool_config.max_keepalive == 5
    assert dispatcher._pool_config.idle_timeout == 30
    assert dispatcher._pool_config.http2 is True


@pytest.mark.unit
@pytest.mark.asyncio
async def test_dispatch_uses_backend_pool_not_shared_client() -> None:
    """dispatch() uses the per-backend pool, not a single shared client."""
    dispatcher = BackendDispatcher()
    # Ensure Core is not loaded so dispatch falls through to backend
    dispatcher._core_engine = None
    backend = _make_backend(base_url="http://mock:11434")
    dispatcher._backends = [backend]

    # Create a mock pool for this backend (OpenAI-compatible format)
    mock_pool = AsyncMock(spec=httpx.AsyncClient)
    mock_response = httpx.Response(
        status_code=200,
        json={
            "id": "chatcmpl-test",
            "object": "chat.completion",
            "created": 1000000,
            "model": "llama3.2",
            "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}],
            "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
        },
        request=httpx.Request("POST", "http://mock:11434/v1/chat/completions"),
    )
    mock_pool.post = AsyncMock(return_value=mock_response)
    dispatcher._pool_registry["http://mock:11434"] = mock_pool

    from llmhosts.proxy.models import UnifiedMessage, UnifiedRequest

    req = UnifiedRequest(
        model="llama3.2",
        messages=[UnifiedMessage(role="user", content="Hello")],
    )
    resp = await dispatcher.dispatch(req)

    mock_pool.post.assert_awaited_once()
    assert resp.content == "Hi"

    # Clean up (don't call shutdown, mock pool has no real aclose)
    dispatcher._pool_registry.clear()


@pytest.mark.unit
@pytest.mark.asyncio
async def test_pool_created_lazily_not_at_startup() -> None:
    """Pools are created on first use, not during startup."""
    dispatcher = BackendDispatcher()
    # Ensure Core is not loaded
    dispatcher._core_engine = None

    await dispatcher.startup()

    # No pools created during startup without Core or discovery
    assert len(dispatcher._pool_registry) == 0
    assert len(dispatcher._backends) == 0

    # Pool created lazily on first use
    backend = _make_backend()
    pool = dispatcher._get_or_create_pool(backend)
    assert pool is not None
    assert "http://127.0.0.1:11434" in dispatcher._pool_registry

    await dispatcher.shutdown()
