#!/usr/bin/env python3
"""Probe each Volcano key against the three plan endpoints to classify it:
  - CODING plan: /api/coding/chat/completions  (ark-code-latest)
  - AGENT plan:  /api/plan/v3/chat/completions (ark-code-latest)
  - PAYGO:       /api/v3/chat/completions      (doubao model name)
Updates vault detail and prints a classification table.
"""
import json, urllib.request, urllib.error
from pathlib import Path

H = Path(__file__).resolve().parent / "llm-key-hunter"
VK = H / "usable_keys" / "volcanoark" / "keys.json"

UA = "Mozilla/5.0 Chrome/124.0 Safari/537.36"


def probe(key, url, model):
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": "hi"}],
                       "max_tokens": 5}).encode()
    req = urllib.request.Request(url, data=body, headers={
        "Authorization": "Bearer " + key, "Content-Type": "application/json",
        "User-Agent": UA})
    try:
        r = urllib.request.urlopen(req, timeout=15)
        return True, f"{r.status}"
    except urllib.error.HTTPError as e:
        msg = e.read()[:140].decode("utf-8", "replace")
        # 404 UnsupportedModel/NotFound means endpoint reachable but plan/model not entitled
        low = msg.lower()
        if e.code in (401, 403):
            return False, f"{e.code} auth"
        if "unsupported" in low or "notsupport" in low or "does not support" in low:
            return False, f"{e.code} no-plan"
        if "not activated" in low or "notopen" in low or "has not activated" in low:
            return False, f"{e.code} no-model"
        if "not found" in low or "notfound" in low:
            return False, f"{e.code} notfound"
        return False, f"{e.code} {msg[:60]}"
    except Exception as e:
        return False, type(e).__name__


def main():
    data = json.loads(VK.read_text())
    for e in data:
        key = e["key"]
        if not key.startswith(("e3e48dbf", "8b7e0c9c", "a6901ee9", "70e5f661", "a0f8cf60")):
            continue
        plans = []
        # PAYGO: try the doubao model that previously returned 200
        paygo_model = e.get("detail", "").split("doubao")[-1].split(" ")[0]
        if paygo_model and not paygo_model[0].isalpha():
            paygo_model = "doubao" + paygo_model
        ok, msg = probe(key, "https://ark.cn-beijing.volces.com/api/v3/chat/completions",
                        paygo_model or "doubao-seed-1-6-250615")
        paygo_ok = ok
        # CODING plan
        c_ok, c_msg = probe(key, "https://ark.cn-beijing.volces.com/api/coding/chat/completions",
                            "ark-code-latest")
        # AGENT plan
        a_ok, a_msg = probe(key, "https://ark.cn-beijing.volces.com/api/plan/v3/chat/completions",
                            "ark-code-latest")
        label = []
        if c_ok: label.append("CODING")
        if a_ok: label.append("AGENT")
        if paygo_ok: label.append("PAYGO")
        plan_str = "+".join(label) if label else "NONE"
        print(f"{key[:14]}  PAYGO={paygo_ok}({msg})  CODING={c_ok}({c_msg})  "
              f"AGENT={a_ok}({a_msg})  => {plan_str}")
        # update vault detail
        e["detail"] = (f"{plan_str} plan; paygo model: {paygo_model}; "
                       f"coding={c_ok} agent={a_ok}")
        e["plans"] = label
    VK.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    print("\nvault updated.")


if __name__ == "__main__":
    main()