"""Phase 5: Dispatcher remote routing — forward requests to remote devices via tunnel URL."""

from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock

import httpx
import pytest

from llmhosts.proxy.dispatcher import BackendDispatcher
from llmhosts.proxy.models import UnifiedMessage, UnifiedRequest

# ---------------------------------------------------------------------------
# Remote backend registration
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_register_remote_backend_normalizes_url_and_registers() -> None:
    """register_remote_backend strips trailing slash and adds a remote BackendEntry."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    dispatcher.register_remote_backend(
        tunnel_url="https://device-abc.llmhosts.com/",
        name="my-remote",
        models=["llama3.2:8b"],
        device_id="dev-1",
    )

    remote = [b for b in dispatcher._backends if b.backend_type == "remote"]
    assert len(remote) == 1
    assert remote[0].name == "my-remote"
    assert remote[0].base_url == "https://device-abc.llmhosts.com"
    assert remote[0].models == ["llama3.2:8b"]
    assert remote[0].device_id == "dev-1"
    assert remote[0].priority == 7

    await dispatcher.shutdown()


@pytest.mark.unit
async def test_register_remote_backend_empty_url_skipped() -> None:
    """register_remote_backend with empty tunnel_url does not add a backend."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    dispatcher.register_remote_backend(
        tunnel_url="",
        name="bad",
        models=[],
    )
    remote = [b for b in dispatcher._backends if b.backend_type == "remote"]
    assert len(remote) == 0

    await dispatcher.shutdown()


@pytest.mark.unit
async def test_register_remote_backends_from_device_dicts() -> None:
    """register_remote_backends builds backends from list of device-like dicts."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    devices: list[dict[str, Any]] = [
        {
            "id": "d1",
            "name": "node-a",
            "tunnel_url": "https://node-a.example.com",
            "models": ["llama3.2:8b", "qwen3:32b"],
        },
        {
            "id": "d2",
            "name": "node-b",
            "tunnel_url": "https://node-b.example.com/",
            "models": '["mistral:7b"]',  # JSON string like Prisma
        },
    ]

    count = dispatcher.register_remote_backends(devices)
    assert count == 2

    remote = [b for b in dispatcher._backends if b.backend_type == "remote"]
    assert len(remote) == 2
    names = {b.name for b in remote}
    assert names == {"node-a", "node-b"}
    by_name = {b.name: b for b in remote}
    assert by_name["node-a"].models == ["llama3.2:8b", "qwen3:32b"]
    assert by_name["node-b"].models == ["mistral:7b"]
    assert by_name["node-b"].base_url == "https://node-b.example.com"

    await dispatcher.shutdown()


@pytest.mark.unit
async def test_unregister_remote_backends_removes_all_remote() -> None:
    """unregister_remote_backends removes only remote backends."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    dispatcher.register_remote_backend(
        tunnel_url="https://r1.example.com",
        name="r1",
        models=["m1"],
    )
    dispatcher.register_remote_backend(
        tunnel_url="https://r2.example.com",
        name="r2",
        models=["m2"],
    )
    assert len([b for b in dispatcher._backends if b.backend_type == "remote"]) == 2

    removed = dispatcher.unregister_remote_backends()
    assert removed == 2
    assert len([b for b in dispatcher._backends if b.backend_type == "remote"]) == 0

    await dispatcher.shutdown()


# ---------------------------------------------------------------------------
# Dispatch to remote (OpenAI-compat forward)
# ---------------------------------------------------------------------------

_OPENAI_COMPAT_RESPONSE = {
    "id": "chatcmpl-remote-1",
    "object": "chat.completion",
    "created": 1234567890,
    "model": "llama3.2:8b",
    "choices": [
        {
            "index": 0,
            "message": {"role": "assistant", "content": "Remote reply"},
            "finish_reason": "stop",
        }
    ],
    "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
}


@pytest.mark.unit
async def test_dispatch_remote_forwards_to_tunnel_url() -> None:
    """Dispatch with remote backend selected forwards POST to tunnel_url/v1/chat/completions."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    tunnel_base = "https://remote-node.example.com"
    dispatcher.register_remote_backend(
        tunnel_url=tunnel_base,
        name="remote-node",
        models=["llama3.2:8b"],
    )

    # Get the pool for this backend and mock POST
    remote_entry = next(b for b in dispatcher._backends if b.backend_type == "remote")
    client = dispatcher._get_or_create_pool(remote_entry)

    async def _fake_post(url: str, **kwargs: Any) -> httpx.Response:
        assert url == f"{tunnel_base}/v1/chat/completions"
        return httpx.Response(
            status_code=200,
            json=_OPENAI_COMPAT_RESPONSE,
            request=httpx.Request("POST", url),
        )

    client.post = AsyncMock(side_effect=_fake_post)  # type: ignore[method-assign]

    req = UnifiedRequest(
        model="llama3.2:8b",
        messages=[UnifiedMessage(role="user", content="Hi")],
        stream=False,
    )
    resp = await dispatcher.dispatch(req)

    assert resp.content == "Remote reply"
    assert resp.model == "llama3.2:8b"
    assert resp.prompt_tokens == 5
    assert resp.completion_tokens == 3
    assert resp.total_tokens == 8
    client.post.assert_called_once()  # type: ignore[attr-defined]

    await dispatcher.shutdown()


class _FakeStreamResponse:
    """Async context manager that yields SSE lines."""

    def __init__(self, lines: list[str]) -> None:
        self._lines = lines
        self.status_code = 200

    async def __aenter__(self) -> _FakeStreamResponse:
        return self

    async def __aexit__(self, *args: object) -> None:
        pass

    def raise_for_status(self) -> None:
        pass

    async def aiter_lines(self) -> Any:
        for line in self._lines:
            yield line


@pytest.mark.unit
async def test_dispatch_stream_remote_forwards_to_tunnel_url() -> None:
    """Streaming dispatch with remote backend uses stream POST to tunnel_url/v1/chat/completions."""
    dispatcher = BackendDispatcher()
    await dispatcher.startup()

    tunnel_base = "https://stream-remote.example.com"
    dispatcher.register_remote_backend(
        tunnel_url=tunnel_base,
        name="stream-remote",
        models=["llama3.2:8b"],
    )

    remote_entry = next(b for b in dispatcher._backends if b.backend_type == "remote")
    client = dispatcher._get_or_create_pool(remote_entry)

    sse_lines = [
        'data: {"choices":[{"delta":{"content":"Hi"},"finish_reason":null}]}',
        'data: {"choices":[{"delta":{"content":"!"},"finish_reason":null}]}',
        'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}',
        "data: [DONE]",
    ]

    def _fake_stream(method: str, url: str, **kwargs: Any) -> _FakeStreamResponse:
        assert method == "POST"
        assert url == f"{tunnel_base}/v1/chat/completions"
        return _FakeStreamResponse(sse_lines)

    client.stream = _fake_stream  # type: ignore[method-assign]

    req = UnifiedRequest(
        model="llama3.2:8b",
        messages=[UnifiedMessage(role="user", content="Hi")],
        stream=True,
    )
    tokens: list[str] = []
    async for token in dispatcher.dispatch_stream(req):
        tokens.append(token)
    assert tokens == ["Hi", "!"]
    await dispatcher.shutdown()
