#!/usr/bin/env python3
"""
usage_collector_kimi_local.py — Kimi 用量采集（Playwright + 本地 CDP）
=====================================================================
连接本地 Chrome CDP (127.0.0.1:9222, chrome-kimi-cdp.service)，
打开 Kimi 控制台，提取用量数据，保存到 usage_stats_kimi.json。

替代旧的 SSH-to-Windows 方案，直接在 246 本地运行。

调用方式: python3 usage_collector_kimi_local.py
"""
import json
import os
import re
import sys
import time
from pathlib import Path
from datetime import datetime, timezone

from playwright.sync_api import sync_playwright

CDP_URL = "http://127.0.0.1:9222"
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 extract_metric(text, keyword, after_window=200):
    """Find percentage and date reset after a keyword in page text."""
    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:
            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:
            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 collect_metrics(page_text):
    """Parse usage metrics from page text."""
    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(page_text, kw)
        if monthly_pct is not None:
            break
    for kw in rolling_keywords:
        rolling_pct, rolling_date, rolling_status, rolling_err = extract_metric(page_text, kw)
        if rolling_pct is not None:
            break
    for kw in weekly_keywords:
        weekly_pct, weekly_date, weekly_status, weekly_err = extract_metric(page_text, kw)
        if weekly_pct is not None:
            break

    # Fallback: positional assignment if keywords miss
    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+)?)%', page_text))
        date_matches = list(re.finditer(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', page_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))

        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"

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


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

    with sync_playwright() as p:
        # Connect to the persistent Chrome CDP service
        try:
            browser = p.chromium.connect_over_cdp(CDP_URL)
        except Exception as e:
            print(f"[kimi] ERROR: Cannot connect to CDP at {CDP_URL}: {e}")
            print("[kimi] Is chrome-kimi-cdp.service running? (systemctl status chrome-kimi-cdp)")
            sys.exit(1)

        # Use the default browser context (has persistent profile with login session)
        context = browser.contexts[0] if browser.contexts else browser.new_context()

        # Check if already logged in by looking for existing kimi pages
        kimi_page = None
        for page in context.pages:
            if "kimi.com" in page.url:
                kimi_page = page
                print(f"[kimi] Reusing existing Kimi tab: {page.url}")
                break

        if not kimi_page:
            kimi_page = context.new_page()

        # Navigate to the quota page
        print(f"[kimi] Navigating to {KIMI_CONSOLE}")
        kimi_page.goto(KIMI_CONSOLE, wait_until="networkidle", timeout=30000)

        # Wait for usage data to render (poll for % signs)
        page_text = ""
        for attempt in range(20):
            page_text = kimi_page.inner_text("body") or ""
            pcts = re.findall(r'(\d+(?:\.\d+)?)%', page_text)
            non_zero = sum(1 for pct in pcts if float(pct) > 0)
            if non_zero >= 2:
                print(f"[kimi] Page ready — {non_zero} non-zero %s (attempt {attempt + 1})")
                break
            time.sleep(1)
        else:
            print(f"[kimi] Warning: page may not have fully loaded")

        # Check for login redirect
        current_url = kimi_page.url
        if "login" in current_url or "passport" in current_url:
            print(f"[kimi] ERROR: Redirected to login page: {current_url}")
            print("[kimi] Need to login first — run: python3 kimi_login.py")
            browser.close()
            sys.exit(1)

        # Parse metrics
        metrics = collect_metrics(page_text)
        print(f"[kimi] Monthly: {metrics['monthly_pct']}% | "
              f"5h: {metrics['rolling_pct']}% | 7d: {metrics['weekly_pct']}%")

        if metrics["errors"]:
            print(f"[kimi] Warnings: {'; '.join(metrics['errors'])}")

        # Screenshot for debugging
        try:
            kimi_page.screenshot(path=str(TEMP_DIR / "kimi_console.png"))
            print("[kimi] Screenshot saved")
        except Exception:
            pass

        # Don't close the page — leave it open for next collection
        # browser.close() would kill the entire CDP service

    # 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()
