#!/usr/bin/env python3
"""
Deep-verify SiliconFlow keys that passed /v1/models.
For each key:
  1. GET /v1/user/info      -> retrieve balance/charge
  2. POST /v1/chat/completions with Qwen/Qwen2.5-7B-Instruct (max_tokens=1)
Classification:
  USABLE      chat returns 200 with choices
  NO_BALANCE  balance=0 / 402 / charge balance exhausted
  NO_ACCESS   403 (realname), 400 (model not granted), 404, 5xx
  UNKNOWN     network/timeout
  DEAD        401 only
"""
import argparse
import json
import urllib.request
import urllib.error
import re
import sys
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, str(Path(__file__).resolve().parent))
from verify_cache import CachedVerifier

BASE = "https://api.siliconflow.cn"
UA = "curl/8.5.0"
OUT = Path("results/medium/SiliconFlow")
OUT.mkdir(parents=True, exist_ok=True)

# Free models — these should work on any identity-verified account with
# even zero balance (SiliconFlow grants free quota to a rotating list).
CHAT_MODEL = "Qwen/Qwen2.5-7B-Instruct"
PINCH_MODEL = "deepseek-ai/DeepSeek-V3"  # paid model, useful balance probe


def http(method, path, key, body=None, timeout=20):
    url = BASE + path
    headers = {
        "Authorization": f"Bearer {key}",
        "User-Agent": UA,
        "Accept": "*/*",
    }
    data = None
    if body is not None:
        data = json.dumps(body).encode()
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.getcode(), resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        try:
            return e.code, e.read().decode("utf-8", "replace")
        except Exception:
            return e.code, ""
    except Exception as e:
        return 0, f"network: {type(e).__name__}: {e}"


def verify(key):
    # Layer 1: user info (balance)
    info_code, info_body = http("GET", "/v1/user/info", key)
    balance = None
    status = None
    if info_code == 200:
        try:
            j = json.loads(info_body)
            data = j.get("data") or {}
            balance = data.get("balance")
            status = data.get("status")
        except Exception:
            pass

    # Layer 2: free model chat
    chat_code, chat_body = http("POST", "/v1/chat/completions", key, {
        "model": CHAT_MODEL,
        "messages": [{"role": "user", "content": "hi"}],
        "max_tokens": 1,
        "temperature": 0,
    })

    # Layer 3: paid model (only if free worked, to detect real balance)
    paid_code, paid_body = 0, ""
    if chat_code == 200:
        paid_code, paid_body = http("POST", "/v1/chat/completions", key, {
            "model": PINCH_MODEL,
            "messages": [{"role": "user", "content": "hi"}],
            "max_tokens": 1,
            "temperature": 0,
        })

    # Classify
    if chat_code == 401:
        return "DEAD", f"401 on chat; info={info_code}"
    if chat_code == 200 and '"choices"' in chat_body:
        if paid_code == 200:
            return "USABLE", f"balance={balance} status={status} free+paid OK"
        return "USABLE_FREE_ONLY", f"balance={balance} status={status} free OK paid={paid_code}"
    if chat_code in (402,):
        return "NO_BALANCE", f"balance={balance}; chat={chat_code} {chat_body[:120]}"
    if chat_code == 403:
        return "NO_ACCESS", f"403: {chat_body[:160]}"
    if chat_code == 400:
        # could be model not allowed, or content filter
        if "balance" in chat_body.lower() or "insufficient" in chat_body.lower():
            return "NO_BALANCE", f"balance={balance}; 400 {chat_body[:160]}"
        return "NO_ACCESS", f"400: {chat_body[:160]}"
    if chat_code == 0:
        return "UNKNOWN", f"network chat: {chat_body[:160]}"
    if 500 <= chat_code < 600:
        return "NO_ACCESS", f"5xx {chat_code}: {chat_body[:120]}"
    return "NO_ACCESS", f"HTTP {chat_code}: {chat_body[:160]}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=20)
    ap.add_argument("--no-cache", action="store_true")
    args = ap.parse_args()

    src = OUT / "usable.txt"
    rows = []
    for line in src.read_text().splitlines():
        if not line.strip() or "|" not in line:
            continue
        parts = line.split("|", 2)
        key = parts[0]
        src_url = parts[1] if len(parts) > 1 else ""
        prev_detail = parts[2] if len(parts) > 2 else ""
        rows.append((key, src_url, prev_detail))
    print(f"loaded {len(rows)} candidate keys", flush=True)

    buckets = {"USABLE": [], "USABLE_FREE_ONLY": [], "NO_BALANCE": [],
               "NO_ACCESS": [], "UNKNOWN": [], "DEAD": []}

    done = 0
    with CachedVerifier("siliconflow", verify, force=args.no_cache,
                        ttl={"USABLE_FREE_ONLY": 30 * 60}) as ver:
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            futs = {pool.submit(ver, k): (k, u, d) for k, u, d in rows}
            for fut in as_completed(futs):
                k, u, d = futs[fut]
                try:
                    v, detail = fut.result()
                except Exception as e:
                    v, detail = "UNKNOWN", f"exc: {e}"
                buckets[v].append((k, u, detail))
                done += 1
                if done % 20 == 0:
                    print(f"  {done}/{len(rows)}  cache: {ver.hits} hits/{ver.live} live", flush=True)
        print(f"cache: {ver.hits} hits, {ver.live} live queries, {len(ver._cache)} cached", flush=True)

    for label, items in buckets.items():
        path = OUT / f"deep_{label.lower()}.txt"
        with path.open("w") as f:
            for k, u, det in items:
                f.write(f"{k}|{u}|{det}\n")
        print(f"{label:18s} {len(items):4d} -> {path.name}")

    # also emit combined all_non_401 for easy newapi ingestion
    keep = []
    for label in ("USABLE", "USABLE_FREE_ONLY", "NO_BALANCE", "NO_ACCESS", "UNKNOWN"):
        for k, u, det in buckets[label]:
            keep.append(f"{k}|{u}|[{label}] {det}")
    with (OUT / "deep_all_non_401.txt").open("w") as f:
        f.write("\n".join(keep) + "\n")
    print(f"kept (non-401): {len(keep)}")


if __name__ == "__main__":
    main()