#!/usr/bin/env python3
"""
usage_collector_kimi.py — Kimi 用量采集（CDP 浏览器自动化）
==========================================================
连接 Chrome CDP proxy (localhost:3456)，打开 Kimi 控制台，
提取用量数据，保存到 usage_stats_kimi.json。

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

CDP = "http://localhost:3456"
KIMI_CONSOLE = "https://www.kimi.com/membership/subscription?tab=quota"

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

OUTPUT_FILE = TEMP_DIR / "usage_stats_kimi.json"
KEY_ID = "key7"
PROVIDER = "kimi"


def cdp_get(path, timeout=10):
    url = f"{CDP}{path}"
    try:
        with urllib.request.urlopen(url, 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):
    url = f"{CDP}{path}"
    try:
        data = body.encode("utf-8") if isinstance(body, str) else body
        req = urllib.request.Request(url, 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):
    """Get document.body.innerText via CDP eval, unwrap JSON envelope."""
    js = "document.body ? document.body.innerText : ''"
    raw = cdp_eval(target_id, js)
    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, min_pct_count=2, timeout=20, interval=1):
    """Poll page text until keywords found or enough non-zero %s, or timeout.

    Returns (text, attempts)."""
    import re as _re
    for attempt in range(int(timeout / interval)):
        text = get_page_text(target_id)
        # Check 1: keywords present
        if keywords and all(k in text for k in keywords):
            print(f"[kimi] Page ready — keywords found (attempt {attempt + 1})")
            return text, attempt + 1
        # Check 2: at least min_pct_count non-zero percentages
        pcts = _re.findall(r'(\d+(?:\.\d+)?)%', text)
        non_zero = sum(1 for p in pcts if float(p) > 0)
        if non_zero >= min_pct_count:
            print(f"[kimi] Page ready — {non_zero} non-zero %s (attempt {attempt + 1})")
            return text, attempt + 1
        time.sleep(interval)
    text = get_page_text(target_id)
    print(f"[kimi] Poll timeout ({timeout}s), using current state")
    return text, int(timeout / interval)


def extract_metric(text, keyword, after_window=200):
    """Find percentage and date reset after a keyword in page text.

    Returns (usage_percent, date_str, status, error_hint).
    """
    import re as _re
    idx = text.find(keyword)
    if idx < 0:
        return None, None, "unknown", f"keyword '{keyword}' not found"
    after = text[idx:idx + after_window]
    m_pct = _re.search(r'(\d+(?:\.\d+)?)%', after)
    m_date = _re.search(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', after)
    pct = float(m_pct.group(1)) if m_pct else None
    date_str = (m_date.group(1) or m_date.group(2)) if m_date else None
    if pct is None:
        return None, None, "parse_error", f"no % found near '{keyword}'"
    status = "ok"
    if pct >= 95:
        status = "rate-limited"
    elif pct >= 80:
        status = "warn"
    return pct, date_str, status, None


def parse_date(dstr):
    """Parse date string to seconds remaining from now."""
    if not dstr:
        return 0
    now = datetime.now(timezone.utc)
    try:
        if '-' in dstr and ':' not in dstr:
            # YYYY-MM-DD
            d = datetime.strptime(dstr, "%Y-%m-%d")
            return max(0, int((d.replace(tzinfo=timezone.utc) - now).total_seconds()))
        elif '-' in dstr and ':' in dstr:
            # MM-DD HH:MM — assume current year
            year = now.year
            d = datetime.strptime(dstr, "%m-%d %H:%M").replace(year=year, tzinfo=timezone.utc)
            if d < now:
                d = d.replace(year=year + 1)
            return max(0, int((d - now).total_seconds()))
    except Exception:
        pass
    return 0


def try_collect(target_id):
    """One round of collection: poll → parse → return metrics dict.

    Returns {monthly_pct, monthly_reset, rolling_pct, rolling_reset,
             weekly_pct, weekly_reset, page_text} or raises.
    """
    import re as _re

    # Poll until page is ready
    text, _attempts = poll_page_ready(target_id, keywords=["%"])

    # Keyword-based extraction (order-independent, survives layout changes)
    # Keywords extracted from actual Kimi page text (2026-07-19)
    monthly_keywords = ["总使用量", "总用量", "月额度", "月"]
    rolling_keywords = ["5 小时用量", "5小时用量", "5小时", "5h", "5 Hour"]
    weekly_keywords = ["7 天用量", "7天用量", "7天", "7d", "7 Day"]

    monthly_pct = monthly_date = monthly_status = monthly_err = None
    rolling_pct = rolling_date = rolling_status = rolling_err = None
    weekly_pct = weekly_date = weekly_status = weekly_err = None

    for kw in monthly_keywords:
        monthly_pct, monthly_date, monthly_status, monthly_err = extract_metric(text, kw)
        if monthly_pct is not None:
            break
    for kw in rolling_keywords:
        rolling_pct, rolling_date, rolling_status, rolling_err = extract_metric(text, kw)
        if rolling_pct is not None:
            break
    for kw in weekly_keywords:
        weekly_pct, weekly_date, weekly_status, weekly_err = extract_metric(text, kw)
        if weekly_pct is not None:
            break

    # Fallback: if any metric still None, extract all percentages and assign positionally
    none_count = sum(1 for x in [monthly_pct, rolling_pct, weekly_pct] if x is None)
    if none_count > 0:
        print(f"[kimi] Keyword matched {3 - none_count}/3 metrics, filling rest positionally")
        pct_matches = list(_re.finditer(r'(\d+(?:\.\d+)?)%', text))
        date_matches = list(_re.finditer(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', text))

        entries = []
        for m in pct_matches:
            pct_val = float(m.group(1))
            pct_pos = m.start()
            best_date = None
            for d in date_matches:
                if d.start() > pct_pos:
                    best_date = d.group(1) or d.group(2)
                    break
            entries.append((pct_val, best_date, pct_pos))

        # Assign in order: monthly=first, rolling=second, weekly=third
        metric_keys = ["monthly", "rolling", "weekly"]
        current = {"monthly": monthly_pct, "rolling": rolling_pct, "weekly": weekly_pct}
        for i, key in enumerate(metric_keys):
            if current[key] is None and i < len(entries):
                if key == "monthly":
                    monthly_pct = entries[i][0]
                    monthly_date = entries[i][1]
                    monthly_status = "ok"
                elif key == "rolling":
                    rolling_pct = entries[i][0]
                    rolling_date = entries[i][1]
                    rolling_status = "ok"
                elif key == "weekly":
                    weekly_pct = entries[i][0]
                    weekly_date = entries[i][1]
                    weekly_status = "ok"

    # Fill None → 0 for output
    monthly_pct = monthly_pct if monthly_pct is not None else 0
    rolling_pct = rolling_pct if rolling_pct is not None else 0
    weekly_pct = weekly_pct if weekly_pct is not None else 0

    return {
        "monthly_pct": monthly_pct,
        "monthly_reset": parse_date(monthly_date),
        "monthly_status": monthly_status if monthly_status else "ok",
        "rolling_pct": rolling_pct,
        "rolling_reset": parse_date(rolling_date),
        "rolling_status": rolling_status if rolling_status else "ok",
        "weekly_pct": weekly_pct,
        "weekly_reset": parse_date(weekly_date),
        "weekly_status": weekly_status if weekly_status else "ok",
        "page_text": text,
        "errors": [e for e in [monthly_err, rolling_err, weekly_err] if e],
    }


def ensure_cdp_proxy():
    """Try to start CDP proxy if it's not running."""
    import subprocess
    result = cdp_get("/targets", timeout=3)
    if isinstance(result, dict) and "error" not in result:
        return True
    if isinstance(result, list):
        return True
    # CDP proxy not running, try to start it
    print("[kimi] CDP proxy not running, attempting to start...")
    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
        )
        print(f"[kimi] CDP start: {proc.stdout.strip()}")
        if proc.returncode == 0:
            # Verify
            result = cdp_get("/targets", timeout=5)
            if isinstance(result, list):
                print("[kimi] CDP proxy started OK")
                return True
    except Exception as e:
        print(f"[kimi] CDP auto-start failed: {e}")
    print("[kimi] ERROR: CDP proxy unavailable and auto-start failed")
    return False


def main():
    print("[kimi] Starting collection...")

    if not ensure_cdp_proxy():
        return

    # 1. Find tab with the membership subscription page
    targets = cdp_get("/targets") or []
    # Guard: CDP proxy may return {"error": ...} dict instead of list
    if isinstance(targets, dict):
        print(f"[kimi] ERROR: CDP proxy returned error: {targets.get('error', 'unknown')}")
        print("[kimi] Is CDP proxy running on localhost:3456? Is Chrome remote debugging enabled?")
        return
    kimi_tab = None
    for t in targets:
        url = t.get("url", "")
        if "membership/subscription" in url or "kimi.com/membership" in url:
            kimi_tab = t
            break

    if kimi_tab:
        target_id = kimi_tab.get("targetId") or kimi_tab.get("id")
        print(f"[kimi] Found quota tab: {target_id}")
        # Always reload for fresh data
        cdp_get(f"/navigate?target={target_id}&url={KIMI_CONSOLE}")
    else:
        print(f"[kimi] No quota tab found, opening new tab...")
        result = cdp_get(f"/new?url={KIMI_CONSOLE}")
        if isinstance(result, dict):
            target_id = result.get("targetId") or result.get("id", "")
        else:
            target_id = str(result)
        print(f"[kimi] New tab: {target_id}")

    if not target_id:
        print("[kimi] ERROR: no target ID")
        return

    # 2. Collect with retry
    metrics = None
    for retry in range(2):
        print(f"[kimi] Collection round {retry + 1}...")
        metrics = try_collect(target_id)
        print(f"[kimi] Monthly: {metrics['monthly_pct']}% | "
              f"5h: {metrics['rolling_pct']}% | 7d: {metrics['weekly_pct']}%")

        # Success criteria: at least one metric has non-zero value
        if metrics["monthly_pct"] > 0 or metrics["rolling_pct"] > 0 or metrics["weekly_pct"] > 0:
            print("[kimi] Non-zero data collected")
            break
        if retry == 0:
            print("[kimi] All zero, reloading and retrying...")
            cdp_get(f"/navigate?target={target_id}&url={KIMI_CONSOLE}")
    else:
        if metrics["errors"]:
            print(f"[kimi] Warnings: {'; '.join(metrics['errors'])}")

    # 3. Take a screenshot for debugging
    try:
        cdp_get(f"/screenshot?target={target_id}&file={TEMP_DIR}/kimi_console.png")
        print("[kimi] Screenshot saved")
    except:
        pass

    # 4. Build output
    output = {
        "ok": True,
        "last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "provider": PROVIDER,
        "accounts": [{
            "key_id": KEY_ID,
            "label": "key7 (Kimi)",
            "provider": PROVIDER,
            "workspace_id": "kimi",
            "rolling": {
                "usage_percent": metrics["rolling_pct"],
                "reset_in_sec": metrics["rolling_reset"],
                "status": metrics["rolling_status"],
            },
            "weekly": {
                "usage_percent": metrics["weekly_pct"],
                "reset_in_sec": metrics["weekly_reset"],
                "status": metrics["weekly_status"],
            },
            "monthly": {
                "usage_percent": metrics["monthly_pct"],
                "reset_in_sec": metrics["monthly_reset"],
                "status": metrics["monthly_status"],
            },
            "session_expired": False,
            "last_update_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "error": None,
        }],
    }

    with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    print(f"[kimi] Output saved to {OUTPUT_FILE}")


if __name__ == "__main__":
    main()