#!/usr/bin/env python3
"""
usage_collector_sensenova.py — SenseNova Token Plan 用量采集（CDP 浏览器自动化）
==========================================================
连接 Chrome CDP proxy (localhost:3456)，打开 SenseNova 控制台，
提取 3 个模型的 5 小时窗口剩余量，保存到 usage_stats_sensenova.json。

控制台（2026-07-23 实测）：
  sensenova-6.7-flash-lite  99.73%剩余
  sensenova-u1-fast         100%剩余
  deepseek-v4-flash         99.80%剩余
公测期免费，各模型 5 小时窗口独立计数，账户下所有 Key 共享配额。

调用方式: python usage_collector_sensenova.py
"""
import json, os, sys, time, urllib.request, re
from pathlib import Path
from datetime import datetime, timezone

CDP = "http://localhost:3456"
CONSOLE_URL = "https://platform.sensenova.cn/console"

SCRIPT_DIR = Path(__file__).resolve().parent
GATEWAY_DIR = SCRIPT_DIR.parent
TEMP_DIR = GATEWAY_DIR / "temp"
os.makedirs(str(TEMP_DIR), exist_ok=True)

# 停止标记：存在此文件时脚本立即退出
STOP_FILE = SCRIPT_DIR / "STOP_SENSNOVA_COLLECT"
if STOP_FILE.exists():
    print(f"[STOPPED] {STOP_FILE.name} exists, exiting.")
    sys.exit(0)

OUTPUT_FILE = TEMP_DIR / "usage_stats_sensenova.json"
PROVIDER = "sensenova"

# 控制台三个模型 → pseudo-key 映射（/api/keys 展示用）
MODEL_KEYS = [
    ("sensenova-6.7-flash-lite", "key8", "SenseNova 6.7 Flash-Lite"),
    ("sensenova-u1-fast", "key9", "SenseNova U1-Fast"),
    ("deepseek-v4-flash", "key10", "SenseNova DeepSeek-V4-Flash"),
]


def cdp_get(path, timeout=10):
    try:
        with urllib.request.urlopen(f"{CDP}{path}", timeout=timeout) as r:
            return json.loads(r.read().decode())
    except Exception as e:
        return {"error": str(e)}


def cdp_post(path, body="", timeout=15):
    try:
        data = body.encode("utf-8") if isinstance(body, str) else body
        req = urllib.request.Request(f"{CDP}{path}", data=data, method="POST")
        req.add_header("Content-Type", "text/plain")
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.read().decode()
    except Exception as e:
        return f"ERROR: {e}"


def cdp_eval(target_id, js, timeout=10):
    return cdp_post(f"/eval?target={target_id}", js, timeout)


def get_page_text(target_id):
    raw = cdp_eval(target_id, "document.body ? document.body.innerText : ''")
    try:
        parsed = json.loads(raw) if isinstance(raw, str) else raw
        return parsed.get("value", "") if isinstance(parsed, dict) else str(raw)
    except (json.JSONDecodeError, AttributeError):
        return str(raw)


def poll_page_ready(target_id, keywords=None, timeout=25, interval=1):
    for attempt in range(int(timeout / interval)):
        text = get_page_text(target_id)
        if keywords and all(k in text for k in keywords):
            print(f"[sensenova] ready (attempt {attempt + 1})")
            return text
        time.sleep(interval)
    print(f"[sensenova] poll timeout, using current state")
    return get_page_text(target_id)


def extract_model_quotas(text):
    """从控制台文本提取各模型剩余百分比。
    实测格式（2026-07-23）：'sensenova-6.7-flash-lite\\n\\n99.73%剩余\\n\\nsensenova-u1-fast\\n\\n100%剩余...'
    注意图表图例区有模型名连写（sensenova-6.7-flash-litesensenova-u1-fast...），
    必须锚定"模型名+换行+数字%剩余"结构，否则会误匹配图例位置抓到别家数值。
    返回 {model_name: remaining_pct}"""
    quotas = {}
    for model, _, _ in MODEL_KEYS:
        m = re.search(re.escape(model) + r'\s*\n+\s*(\d+(?:\.\d+)?)\s*%\s*剩余', text)
        if m:
            quotas[model] = float(m.group(1))
    return quotas


def ensure_cdp_proxy():
    result = cdp_get("/targets", timeout=3)
    if isinstance(result, list):
        return True
    print("[sensenova] CDP proxy not running, attempting to start...")
    import subprocess
    ps_script = r"D:\F\NewI\opencode\daily-workspace\.opencode\scripts\start-cdp-proxy.ps1"
    try:
        proc = subprocess.run(
            ["powershell", "-ExecutionPolicy", "Bypass", "-File", ps_script],
            capture_output=True, text=True, timeout=20)
        if proc.returncode == 0 and isinstance(cdp_get("/targets", timeout=5), list):
            print("[sensenova] CDP proxy started OK")
            return True
    except Exception as e:
        print(f"[sensenova] CDP auto-start failed: {e}")
    return False


def find_or_open_console_tab():
    targets = cdp_get("/targets")
    if not isinstance(targets, list):
        return None
    for t in targets:
        if "platform.sensenova.cn" in (t.get("url") or ""):
            return t.get("targetId") or t.get("id")
    # open new background tab
    r = cdp_get(f"/new?url={CONSOLE_URL}", timeout=30)
    if isinstance(r, dict) and r.get("targetId"):
        return r["targetId"]
    print(f"[sensenova] open tab failed: {r}")
    return None


def main():
    print("[sensenova] Starting collection...")
    if not ensure_cdp_proxy():
        print("[sensenova] ERROR: CDP proxy unavailable")
        sys.exit(1)

    tid = find_or_open_console_tab()
    if not tid:
        print("[sensenova] ERROR: cannot get console tab")
        sys.exit(1)

    text = poll_page_ready(tid, keywords=["剩余"])
    quotas = extract_model_quotas(text)
    print(f"[sensenova] quotas: {quotas}")

    now_iso = datetime.now(timezone.utc).isoformat()
    accounts = []
    for model, key_id, label in MODEL_KEYS:
        remaining = quotas.get(model)
        if remaining is None:
            accounts.append({
                "key_id": key_id, "label": label, "provider": PROVIDER,
                "last_update_iso": now_iso,
                "error": "quota not found on console page",
                "rolling": {"usage_percent": None, "reset_in_sec": None, "status": "unknown"},
            })
            continue
        usage = round(100.0 - remaining, 2)
        status = "ok" if usage < 80 else ("warn" if usage < 95 else "rate-limited")
        accounts.append({
            "key_id": key_id, "label": label, "provider": PROVIDER,
            "last_update_iso": now_iso,
            "error": None,
            "rolling": {"usage_percent": usage, "reset_in_sec": 5 * 3600, "status": status},
            "monthly": None, "weekly": None,
            "remaining_percent": remaining,
            "note": "公测免费·5小时窗口·账户所有Key共享",
        })

    out = {
        "source": "usage_collector_sensenova.py (CDP scrape of platform.sensenova.cn/console)",
        "collected_at": now_iso,
        "accounts": accounts,
    }
    with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f:
        json.dump(out, f, ensure_ascii=False, indent=2)
    print(f"[sensenova] saved → {OUTPUT_FILE} ({sum(1 for a in accounts if not a.get('error'))}/3 models)")


if __name__ == "__main__":
    main()