#!/usr/bin/env python3
"""
NewAPI 模型健康巡检（快速版）
每 6h 运行，测试关键模型的响应状态
输出: ~/.hermes/model-health.json
"""

import json
import time
import urllib.request
import urllib.error
import os
import sys
from datetime import datetime, timezone

API = "http://127.0.0.1:3000/v1"
KEY = "0ExNiLblJvIWBDpkS50fwOBw4MmqLyKdHJK5iQtlw9dOMWBP"
OUTPUT = os.path.expanduser("~/.hermes/model-health.json")

# 测试模型列表（按优先级排列）
# 第1批: 已知候选（先测快的）
BATCH_1 = [
    "minimaxai/minimax-m3",
    "minimaxai/minimax-m2.7",
    "stepfun-ai/step-3.5-flash",
    "deepseek-ai/deepseek-v3.2",
    "microsoft/phi-4-mini-instruct",
    "meta/llama-4-maverick-17b-128e-instruct",
]

# 第2批: 大型模型
BATCH_2 = [
    "mistralai/mistral-medium-3.5-128b",
    "meta/llama-3.3-70b-instruct",
    "qwen/qwen3.5-122b-a10b",
    "z-ai/glm4.7",
    "z-ai/glm5",
    "bytedance/seed-oss-36b-instruct",
]

# 第3批: 超大/专用
BATCH_3 = [
    "mistralai/mistral-large-3-675b-instruct-2512",
    "qwen/qwen3-coder-480b-a35b-instruct",
    "qwen/qwen2.5-coder-32b-instruct",
    "moonshotai/kimi-k2-instruct",
]

ALL_MODELS = BATCH_1 + BATCH_2 + BATCH_3

HEADERS = {
    "Authorization": f"Bearer {KEY}",
    "Content-Type": "application/json",
}

PROMPT = "你好"


def test_model(model: str) -> dict:
    """测试单个模型 2 次，返回汇总"""
    trials = []

    for t in range(2):
        payload = json.dumps({
            "model": model,
            "messages": [{"role": "user", "content": PROMPT}],
            "max_tokens": 20,
        }).encode()

        req = urllib.request.Request(
            f"{API}/chat/completions",
            data=payload,
            headers=HEADERS,
            method="POST",
        )

        start = time.time()
        try:
            with urllib.request.urlopen(req, timeout=15) as resp:
                body = json.loads(resp.read())
        except urllib.error.HTTPError as e:
            trials.append({"status": "fail", "error": f"HTTP_{e.code}", "latency_ms": round((time.time() - start) * 1000)})
            continue
        except Exception as e:
            trials.append({"status": "fail", "error": str(e)[:60], "latency_ms": round((time.time() - start) * 1000)})
            continue

        elapsed = round((time.time() - start) * 1000)

        try:
            choice = body["choices"][0]
            msg = choice.get("message", {})
            content = msg.get("content", "") or ""
            finish = choice.get("finish_reason", "")
            usage = body.get("usage", {})

            # ttft 从 nvext 取，没有就估计
            ttft = body.get("nvext", {}).get("timing", {}).get("ttft_ms", -1)
            if ttft < 0:
                ttft = round(elapsed * 0.3)

            trials.append({
                "status": "ok",
                "latency_ms": elapsed,
                "ttft_ms": ttft,
                "has_content": 1 if content.strip() else 0,
                "completion_tokens": usage.get("completion_tokens", 0),
                "finish_reason": finish,
            })
        except (KeyError, IndexError, json.JSONDecodeError) as e:
            trials.append({"status": "fail", "error": f"parse: {e}", "latency_ms": elapsed})

    # 汇总
    ok_count = sum(1 for t in trials if t["status"] == "ok")
    fail_count = 2 - ok_count

    if ok_count == 2:
        stability = "stable"
    elif ok_count == 1:
        stability = "unstable"
    else:
        stability = "dead"

    ok_trials = [t for t in trials if t["status"] == "ok"]
    avg_latency = round(sum(t["latency_ms"] for t in ok_trials) / len(ok_trials)) if ok_trials else 0
    avg_ttft = round(sum(t.get("ttft_ms", 0) for t in ok_trials) / len(ok_trials)) if ok_trials else -1

    last_ok = ok_trials[-1] if ok_trials else trials[-1]
    last_finish = last_ok.get("finish_reason", "error")

    return {
        "model": model,
        "tests": 2,
        "success": ok_count,
        "failure": fail_count,
        "avg_latency_ms": avg_latency,
        "avg_ttft_ms": avg_ttft,
        "stability": stability,
        "last_status": "ok" if ok_count > 0 else "fail",
        "last_finish": last_finish,
    }


def main():
    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    results = []
    deadline = time.time() + 240  # 4分钟全局超时

    for model in ALL_MODELS:
        if time.time() > deadline:
            print(f"⏰ 全局超时，跳过剩余模型")
            break

        entry = test_model(model)
        results.append(entry)

        icon = "✅" if entry["stability"] == "stable" else ("⚠️" if entry["stability"] == "unstable" else "❌")
        print(f"{icon} {model:45s} {entry['avg_latency_ms']:>6}ms | {entry['success']}/2 ok | {entry['stability']}")
        sys.stdout.flush()

    # 汇总
    healthy = sum(1 for r in results if r["stability"] == "stable")
    flaky = sum(1 for r in results if r["stability"] == "unstable")
    dead = sum(1 for r in results if r["stability"] == "dead")

    stable_models = [r for r in results if r["stability"] == "stable"]
    stable_sorted = sorted(stable_models, key=lambda x: x["avg_latency_ms"])
    fastest = stable_sorted[:5] if stable_sorted else []

    summary = {
        "timestamp": timestamp,
        "total_models": len(results),
        "stable": healthy,
        "unstable": flaky,
        "dead": dead,
        "fastest_stable": [m["model"] for m in fastest],
        "recommendations": {
            "fast": [m["model"] for m in fastest],
            "fastest3": [m["model"] for m in fastest[:3]],
            "priorities": {
                "日常快速": fastest[:3] if len(fastest) >= 3 else fastest,
                "复杂推理": [m["model"] for m in sorted(stable_models, key=lambda x: -x.get("completion_tokens", 0) if hasattr(x, "get") else 0)[:2]],
            },
        },
        "models": results,
    }

    os.makedirs(os.path.dirname(OUTPUT), exist_ok=True)
    with open(OUTPUT + ".new", "w") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    os.replace(OUTPUT + ".new", OUTPUT)

    print(f"\n{'='*50}")
    print(f"巡检完成: {healthy}个稳定 / {flaky}个不稳定 / {dead}个死 (共{len(results)}个)")
    if fastest:
        print(f"推荐: {', '.join(summary['recommendations']['fast'])}")


if __name__ == "__main__":
    main()