#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""验收 LiteLLM：/v1/models、chat、embeddings。"""
from __future__ import annotations

import json
import sys
import urllib.error
import urllib.request

import _bootstrap  # noqa: F401
import _config as LC


def _req(method: str, path: str, body: dict | None = None) -> tuple[int, str]:
    url = LC.LITELLM_BASE.rstrip("/") + path
    data = None if body is None else json.dumps(body).encode("utf-8")
    headers = {
        "Authorization": f"Bearer {LC.LITELLM_MASTER_KEY}",
        "Content-Type": "application/json",
    }
    r = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(r, timeout=180) as resp:
            return resp.status, resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode("utf-8", "replace")
    except Exception as e:  # noqa: BLE001
        return 0, str(e)


def main() -> int:
    print("[INFO] base:", LC.LITELLM_BASE)
    code, text = _req("GET", "/v1/models")
    print(f"[models] HTTP {code}")
    print(text[:500])
    if code != 200:
        print("[FAIL] /v1/models")
        return 1
    models = json.loads(text)
    ids = {m.get("id") for m in models.get("data") or []}
    if LC.LITELLM_CHAT_MODEL not in ids and "qwen" not in ids:
        print("[FAIL] chat model not listed:", ids)
        return 1

    code, text = _req(
        "POST",
        "/v1/chat/completions",
        {
            "model": LC.LITELLM_CHAT_MODEL,
            "messages": [{"role": "user", "content": "用五个字打个招呼"}],
            "max_tokens": 128,
            "temperature": 0.2,
        },
    )
    print(f"[chat] HTTP {code}")
    print(text[:800])
    if code != 200:
        print("[FAIL] chat")
        return 1
    chat = json.loads(text)
    msg = ((chat.get("choices") or [{}])[0].get("message") or {})
    content = msg.get("content") or msg.get("reasoning_content") or ""
    if not str(content).strip():
        print("[FAIL] empty chat content")
        return 1
    print("[OK] chat:", str(content)[:120].replace("\n", " "))

    code, text = _req(
        "POST",
        "/v1/embeddings",
        {"model": LC.LITELLM_EMBED_MODEL, "input": "请假流程测试"},
    )
    print(f"[embeddings] HTTP {code}")
    print(text[:400])
    if code != 200:
        print("[FAIL] embeddings")
        return 1
    emb = json.loads(text)
    vec = ((emb.get("data") or [{}])[0].get("embedding")) or []
    if len(vec) < 8:
        print("[FAIL] embedding vector too short:", len(vec))
        return 1
    print("[OK] embeddings dim=", len(vec))
    print("[PASS] LiteLLM 验收通过")
    return 0


if __name__ == "__main__":
    sys.exit(main())