"""Integration tests for LLMHosts Core dispatcher wiring."""

from __future__ import annotations

import json as _json

import pytest


@pytest.mark.unit
class TestCoreDispatcherSelection:
    def test_core_engine_loaded_state(self) -> None:
        """When Core engine is loaded, it is the default local backend."""
        from llmhosts.proxy.dispatcher import BackendDispatcher

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        mock_engine = type("MockEngine", (), {"is_loaded": True})()
        dispatcher._core_engine = mock_engine
        assert dispatcher._core_engine is not None
        assert dispatcher._core_engine.is_loaded is True

    def test_core_engine_not_loaded(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        mock_engine = type("MockEngine", (), {"is_loaded": False})()
        dispatcher._core_engine = mock_engine
        assert dispatcher._core_engine.is_loaded is False

    def test_core_engine_unavailable(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        dispatcher._core_engine = None
        assert dispatcher._core_engine is None


@pytest.mark.unit
class TestCoreDispatchResponse:
    @pytest.mark.asyncio
    async def test_response_matches_openai_spec(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        class MockEngine:
            is_loaded = True
            model_name = "qwen2.5-0.5b-q4_k_m"
            _tokenizer = None

            def generate(self, prompt: str, **kwargs) -> dict:
                return {"text": "Hello back!", "tokens": [1, 2, 3], "token_count": 3}

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        dispatcher._core_engine = MockEngine()
        dispatcher._config = type("Config", (), {"privacy_mode": None})()

        messages = [{"role": "user", "content": "Hello"}]
        body, headers = await dispatcher._dispatch_core(messages, model="llmhosts")

        assert body["object"] == "chat.completion"
        assert body["id"].startswith("chatcmpl-")
        assert body["model"] == "llmhosts"
        assert len(body["choices"]) == 1
        assert body["choices"][0]["message"]["role"] == "assistant"
        assert body["choices"][0]["message"]["content"] == "Hello back!"
        assert body["choices"][0]["finish_reason"] == "stop"
        assert "usage" in body
        assert body["usage"]["completion_tokens"] == 3

        assert headers["X-LLMHosts-Backend"] == "llmhosts-core"
        assert headers["X-LLMHosts-Cost"] == "0.0"

    @pytest.mark.asyncio
    async def test_privacy_strict_mode_header(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        class MockEngine:
            is_loaded = True
            _tokenizer = None

            def generate(self, prompt: str, **kwargs) -> dict:
                return {"text": "response", "tokens": [1], "token_count": 1}

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        dispatcher._core_engine = MockEngine()
        dispatcher._config = type("Config", (), {"privacy_mode": "strict"})()

        messages = [{"role": "user", "content": "Hello"}]
        _body, headers = await dispatcher._dispatch_core(messages, model="llmhosts")
        assert headers["X-LLMHosts-Privacy"] == "strict"

    @pytest.mark.asyncio
    async def test_stop_sequences_truncate_output(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        class MockEngine:
            is_loaded = True
            _tokenizer = None

            def generate(self, prompt: str, **kwargs) -> dict:
                return {"text": "Hello world STOP rest of text", "tokens": [1, 2, 3, 4, 5], "token_count": 5}

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        dispatcher._core_engine = MockEngine()
        dispatcher._config = type("Config", (), {"privacy_mode": None})()

        messages = [{"role": "user", "content": "Hello"}]
        request = {"stop": ["STOP"]}
        body, _headers = await dispatcher._dispatch_core(messages, model="llmhosts", request=request)
        assert body["choices"][0]["message"]["content"] == "Hello world "

    @pytest.mark.asyncio
    async def test_stream_yields_sse_chunks(self) -> None:
        from llmhosts.proxy.dispatcher import BackendDispatcher

        class MockEngine:
            is_loaded = True
            _tokenizer = None

            def generate(self, prompt: str, **kwargs) -> dict:
                return {"text": "Hello world", "tokens": [1, 2], "token_count": 2}

        dispatcher = BackendDispatcher.__new__(BackendDispatcher)
        dispatcher._core_engine = MockEngine()
        dispatcher._config = type("Config", (), {"privacy_mode": None})()

        messages = [{"role": "user", "content": "Hi"}]
        chunks = []
        async for chunk in dispatcher._dispatch_core_stream(messages, model="llmhosts"):
            chunks.append(chunk)

        assert len(chunks) == 4  # "Hello ", "world", finish_reason, [DONE]
        assert chunks[-1] == "data: [DONE]\n\n"

        first_data = _json.loads(chunks[0].removeprefix("data: ").strip())
        assert first_data["object"] == "chat.completion.chunk"
        assert first_data["choices"][0]["delta"]["content"].startswith("Hello")
