"""模型上下文探测与有效 token 预算。"""

from __future__ import annotations

import asyncio
from typing import Any
from unittest.mock import patch

import httpx

from aicd.agent.llm import OpenAICompatibleClient
from aicd.agent.model_context_probe import probe_model_context_tokens, _scan_context_int
from aicd.config import (
    AgentConfig,
    AppConfig,
    LLMConfig,
    normalize_model_context_tokens_map,
    resolve_llm_model_context_window_tokens,
)


def test_scan_context_int_nested() -> None:
    assert _scan_context_int({"foo": {"context_length": 8192}}) == 8192
    assert _scan_context_int({"model_info": {"general.context_length": 16384}}) == 16384


class _FakeAsyncClient:
    """替换 ``httpx.AsyncClient``：在 ``model_context_probe`` 内仅用到 ``post`` / ``get``。"""

    def __init__(
        self,
        *,
        post_r: httpx.Response | None = None,
        get_r: httpx.Response | list[httpx.Response] | None = None,
    ):
        self._post_r = post_r
        if isinstance(get_r, list):
            self._get_rs = list(get_r)
        elif get_r is None:
            self._get_rs = []
        else:
            self._get_rs = [get_r]

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

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

    async def post(self, url: str, json: dict | None = None) -> httpx.Response:
        assert self._post_r is not None
        return self._post_r

    async def get(self, url: str) -> httpx.Response:
        assert self._get_rs, "no mocked GET response"
        return self._get_rs.pop(0)


def test_probe_ollama_show_parses_context() -> None:
    payload = {"model_info": {"general.context_length": 32768}}
    fake = _FakeAsyncClient(
        post_r=httpx.Response(200, json=payload),
    )
    with patch("aicd.agent.model_context_probe.httpx.AsyncClient", lambda **kw: fake):

        async def _go() -> None:
            toks, src, detail = await probe_model_context_tokens(
                base_url="http://127.0.0.1:11434/v1",
                api_key="x",
                model_id="llama3",
                proto="ollama",
                timeout_sec=5.0,
            )
            assert toks == 32768
            assert "ollama" in src
            assert detail == ""

        asyncio.run(_go())


def test_probe_openai_models_list() -> None:
    data = {"data": [{"id": "my-model", "context_length": 128_000}]}
    fake = _FakeAsyncClient(
        get_r=[
            httpx.Response(404),
            httpx.Response(200, json=data),
        ],
    )
    with patch("aicd.agent.model_context_probe.httpx.AsyncClient", lambda **kw: fake):

        async def _go() -> None:
            toks, src, detail = await probe_model_context_tokens(
                base_url="http://fake/v1",
                api_key="sk-test",
                model_id="my-model",
                proto="openai",
                timeout_sec=5.0,
            )
            assert toks == 128_000
            assert src == "openai_compatible_models_list"
            assert detail == ""

        asyncio.run(_go())


def test_effective_context_budget_tokens() -> None:
    cfg = AppConfig()
    cfg.agent = AgentConfig(context_budget_tokens=500_000)
    cfg.llm = LLMConfig(model="x", max_output_tokens=4096)
    client = OpenAICompatibleClient(cfg)
    client._api_model_context_tokens = 32_000
    eff = client.effective_context_budget_tokens(500_000)
    assert eff == 32_000 - 4096
    client._api_model_context_tokens = None
    assert client.effective_context_budget_tokens(500_000) == 500_000


def test_model_context_tokens_map_and_merge() -> None:
    assert normalize_model_context_tokens_map({"a": 100, "b": -1, "c": "x"}) == {
        "a": 100
    }
    llm = LLMConfig(
        model="m1",
        model_context_tokens={"m1": 16_000, "default": 8_000},
    )
    assert resolve_llm_model_context_window_tokens(llm, "m1") == 16_000
    assert resolve_llm_model_context_window_tokens(llm, "unknown") == 8_000
    cfg = AppConfig()
    cfg.agent = AgentConfig(context_budget_tokens=500_000)
    cfg.llm = LLMConfig(model="x", max_output_tokens=4096)
    cfg.llm.model_context_tokens = {"x": 20_000}
    client = OpenAICompatibleClient(cfg)
    client._api_model_context_tokens = 128_000
    assert client._effective_model_context_window_tokens() == 20_000
    assert client.effective_context_budget_tokens(500_000) == 20_000 - 4096
