"""向 OpenAI 兼容或 Ollama 服务端探测模型上下文窗口，用于裁剪对话历史预算。"""

from __future__ import annotations

import json
from typing import Any
from urllib.parse import quote

import httpx

from aicd.config import PROTO_OLLAMA

_CONTEXT_INT_KEYS = frozenset(
    {
        "context_length",
        "context_window",
        "max_model_len",
        "max_sequence_length",
        "max_sequence_len",
        "n_ctx",
        "model_max_length",
        "max_context_length",
        "max_position_embeddings",
    }
)


def _scan_context_int(obj: Any, *, depth: int = 0) -> int | None:
    """在 JSON 对象中递归查找疑似「上下文 token 上限」的整数字段。"""
    if depth > 8:
        return None
    if isinstance(obj, dict):
        for k, v in obj.items():
            lk = str(k).lower().replace("-", "_")
            if (
                lk in _CONTEXT_INT_KEYS
                or lk.endswith(".context_length")
                or lk.endswith("_context_length")
            ):
                if isinstance(v, int) and v >= 1024:
                    return v
                if isinstance(v, float) and v >= 1024:
                    return int(v)
        for v in obj.values():
            hit = _scan_context_int(v, depth=depth + 1)
            if hit is not None:
                return hit
    elif isinstance(obj, list):
        for it in obj:
            hit = _scan_context_int(it, depth=depth + 1)
            if hit is not None:
                return hit
    return None


def _ollama_root_from_openai_base(base_url: str) -> str | None:
    b = base_url.rstrip("/")
    if b.endswith("/v1"):
        root = b[:-3].rstrip("/")
        return root or None
    return None


async def _probe_ollama_show(
    *,
    openai_base_url: str,
    model_id: str,
    timeout_sec: float,
) -> tuple[int | None, str]:
    root = _ollama_root_from_openai_base(openai_base_url)
    if not root:
        return None, ""
    name = (model_id or "").strip()
    if not name:
        return None, ""
    url = f"{root}/api/show"
    try:
        async with httpx.AsyncClient(timeout=timeout_sec) as ac:
            r = await ac.post(url, json={"name": name})
            if r.status_code != 200:
                return None, f"ollama_show_http_{r.status_code}"
            data = r.json()
    except (OSError, httpx.HTTPError, json.JSONDecodeError, TypeError):
        return None, "ollama_show_error"
    if not isinstance(data, dict):
        return None, "ollama_show_bad_json"
    # 顶层或 model_info 中的 *.context_length
    top = _scan_context_int(data)
    if top is not None:
        return top, "ollama_api_show"
    mi = data.get("model_info")
    if isinstance(mi, dict):
        for k, v in mi.items():
            if isinstance(k, str) and "context_length" in k.lower():
                if isinstance(v, int) and v >= 1024:
                    return v, "ollama_api_show_model_info"
                if isinstance(v, float) and v >= 1024:
                    return int(v), "ollama_api_show_model_info"
        hit = _scan_context_int(mi)
        if hit is not None:
            return hit, "ollama_api_show_model_info"
    return None, "ollama_show_no_context"


async def _probe_openai_models_route(
    *,
    base_url: str,
    api_key: str,
    model_id: str,
    timeout_sec: float,
) -> tuple[int | None, str]:
    b = base_url.rstrip("/")
    headers: dict[str, str] = {}
    if (api_key or "").strip():
        headers["Authorization"] = f"Bearer {api_key.strip()}"

    safe_model = quote(model_id, safe="")
    urls = (
        f"{b}/models/{safe_model}",
        f"{b}/models",
    )

    async with httpx.AsyncClient(timeout=timeout_sec, headers=headers) as ac:
        for u in urls:
            try:
                r = await ac.get(u)
            except (OSError, httpx.HTTPError):
                continue
            if r.status_code != 200:
                continue
            try:
                data = r.json()
            except (json.JSONDecodeError, TypeError):
                continue
            if u.endswith("/models") and isinstance(data, dict):
                rows = data.get("data")
                if isinstance(rows, list):
                    for row in rows:
                        if not isinstance(row, dict):
                            continue
                        rid = str(row.get("id", ""))
                        if rid != model_id:
                            continue
                        hit = _scan_context_int(row)
                        if hit is not None:
                            return hit, "openai_compatible_models_list"
            hit = _scan_context_int(data)
            if hit is not None:
                return hit, "openai_compatible_models_get"
    return None, "models_route_miss"


async def probe_model_context_tokens(
    *,
    base_url: str,
    api_key: str,
    model_id: str,
    proto: str,
    timeout_sec: float = 12.0,
) -> tuple[int | None, str, str]:
    """
    返回 ``(tokens | None, source_tag, detail)``。

    ``source_tag`` 表示成功路径；失败时 ``tokens`` 为 ``None``，``detail`` 为简短原因（可记入日志）。
    """
    mid = (model_id or "").strip()
    if not mid:
        return None, "", "empty_model_id"
    base = (base_url or "").strip().rstrip("/")
    if not base:
        return None, "", "empty_base_url"

    p = (proto or "").strip().lower()
    if p == PROTO_OLLAMA or "11434" in base:
        toks, src = await _probe_ollama_show(
            openai_base_url=base,
            model_id=mid,
            timeout_sec=timeout_sec,
        )
        if toks is not None:
            return toks, src, ""
        # 回退：部分 Ollama 仍挂 /v1/models
        toks2, src2 = await _probe_openai_models_route(
            base_url=base,
            api_key=api_key,
            model_id=mid,
            timeout_sec=timeout_sec,
        )
        if toks2 is not None:
            return toks2, src2, ""
        return None, "", src or "ollama_fallback_miss"

    toks3, src3 = await _probe_openai_models_route(
        base_url=base,
        api_key=api_key,
        model_id=mid,
        timeout_sec=timeout_sec,
    )
    if toks3 is not None:
        return toks3, src3, ""
    return None, "", src3
