"""Tests for the resolved-model stamping on LLMChunk/LLMResponse.

The fallback backend resolves a single model from the profile's priority list
(trying each until one succeeds) and must surface which model actually served
the call so the UI can show it. We verify the resolved model is stamped onto the
first streaming chunk and onto the non-streaming response.
"""

import asyncio

from navi.llm.base import LLMChunk, LLMResponse
from navi.llm.fallback import FallbackOllamaBackend, ServerEntry


def _make_backend(monkeypatch, streaming_gen, complete_resp):
    """Build a FallbackOllamaBackend whose inner OllamaBackend is faked.

    The fake records the model it was asked to use and returns ``streaming_gen``
    / ``complete_resp``. This lets us assert the fallback stamps the resolved
    model that survived the priority list.
    """
    import navi.llm.fallback as fallback_mod

    fallback_mod.clear_blacklists()

    used_models: list[str] = []

    class FakeOllamaBackend:
        def __init__(self, **kwargs):
            self.kwargs = kwargs

        async def complete(self, *args, model=None, **kwargs):
            used_models.append(model)
            if isinstance(complete_resp, Exception):
                raise complete_resp
            return complete_resp

        async def stream_complete(self, *args, model=None, **kwargs):
            used_models.append(model)
            if isinstance(streaming_gen, Exception):
                raise streaming_gen
            for chunk in streaming_gen:
                yield chunk

    monkeypatch.setattr(fallback_mod, "OllamaBackend", FakeOllamaBackend)
    backend = FallbackOllamaBackend([ServerEntry(host="http://ollama.test")])
    return backend, used_models


def test_fallback_stamps_resolved_model_on_response(monkeypatch):
    """complete() returns an LLMResponse carrying the resolved model name."""
    resp = LLMResponse(content="ok", tool_calls=None, finish_reason="stop")
    backend, used = _make_backend(monkeypatch, streaming_gen=[], complete_resp=resp)

    result = asyncio.run(
        backend.complete([], model=["first-choice", "second-choice"])
    )

    assert used == ["first-choice"]
    assert result.model == "first-choice"


def test_fallback_falls_back_and_stamps_serving_model(monkeypatch):
    """When the first model is not found, the second serves and is stamped."""
    import navi.llm.fallback as fallback_mod

    fallback_mod.clear_blacklists()

    served: list[str] = []

    class FakeOllamaBackend:
        def __init__(self, **kwargs):
            pass

        async def complete(self, *args, model=None, **kwargs):
            served.append(model)
            if model == "down-model":
                raise fallback_mod.LLMModelNotFoundError("missing")
            return LLMResponse(content="ok", tool_calls=None, finish_reason="stop")

    monkeypatch.setattr(fallback_mod, "OllamaBackend", FakeOllamaBackend)
    backend = FallbackOllamaBackend([ServerEntry(host="http://ollama.test")])

    result = asyncio.run(
        backend.complete([], model=["down-model", "alive-model"])
    )

    assert served == ["down-model", "alive-model"]
    assert result.model == "alive-model"


def test_fallback_stamps_resolved_model_on_first_chunk(monkeypatch):
    """stream_complete() stamps the resolved model onto the first chunk only."""
    chunks = [
        LLMChunk(delta="hello"),
        LLMChunk(delta=" world"),
        LLMChunk(delta=None, finish_reason="stop", prompt_tokens=5, completion_tokens=2),
    ]
    backend, used = _make_backend(monkeypatch, streaming_gen=chunks, complete_resp=None)

    out = []
    async def collect():
        async for chunk in backend.stream_complete([], model=["the-model"]):
            out.append(chunk)
    asyncio.run(collect())

    assert used == ["the-model"]
    assert len(out) == 3
    assert out[0].model == "the-model"
    assert out[1].model is None
    assert out[2].model is None


def test_ollama_backend_stamps_first_chunk_only(monkeypatch):
    """OllamaBackend stamps the resolved model on the first chunk, not later ones."""
    from navi.llm.ollama import OllamaBackend

    class FakeChunk:
        def __init__(self, content, done=False):
            self.message = type("M", (), {"content": content, "thinking": None, "tool_calls": None})()
            self.done = done
            self.prompt_eval_count = 5 if done else None
            self.eval_count = 2 if done else None

    class FakeClient:
        async def chat(self, **kwargs):
            async def _gen():
                yield FakeChunk("hi")
                yield FakeChunk(" there", done=True)
            return _gen()

    backend = OllamaBackend(model="default-model", host="http://ollama.test", api_key="")
    monkeypatch.setattr(backend, "_client", FakeClient())

    out = []
    async def collect():
        async for chunk in backend.stream_complete([], model=["configured-model"]):
            out.append(chunk)
    asyncio.run(collect())

    assert out[0].model == "configured-model"
    assert out[1].model is None