#!/usr/bin/env python3
"""
usage_collector.py — OpenCode Go 订阅用量采集脚本 (Hybrid v3.0)
================================================================
混合采集模式：
  主模式 — HTTP + stored cookies (全部 4 账号同时采集)
  辅助模式 — CDP (仅当前 Chrome 登录的账号，用于补充/验证)

采集流程:
  1. 遍历 accounts.json 中所有配置了 workspace_id 的账号
  2. 对每个账号：
     a. 主模式: 读取 cookies/{key_id}.json → 构建 Cookie header →
        HTTP GET https://opencode.ai/workspace/{ws_id}/go → 解析 SSR HTML
     b. 辅助模式(仅对 Chrome 当前登录账号): CDP fetch(location.href) → 解析 SSR HTML
  3. 合并结果到 gateway/temp/usage_stats.json

Cookie 提取流程 (一次性，每个账号登录后执行):
  python extract_cookies.py key1 --workspace-id wrk_XXX
  → cookies/key1.json (含 httpOnly auth cookie)

用法:
  python usage_collector.py            # 一次性采集 (全部账号)
  python usage_collector.py --daemon   # 后台守护，每 5 分钟采集一次
  python usage_collector.py --print    # 采集 + 打印结果
  python usage_collector.py --cdp-only # 仅 CDP 采集当前登录账号

相关文件:
  - usage_monitor/accounts.json — 账号配置
  - usage_monitor/cookies/{key_id}.json — 存储的 cookies (gitignored)
  - usage_monitor/extract_cookies.py — Cookie 提取工具
  - gateway/temp/usage_stats.json — 采集结果缓存
  - xmpp_agent_core.py /usage endpoint 触发本脚本
  - CDP proxy (cdp-proxy.mjs) localhost:3456
"""
import os
import sys
import re
import json
import time
import logging
import argparse
import urllib.request
import urllib.error
import urllib.parse
import http.cookiejar
from pathlib import Path
from datetime import datetime, timezone

# ── Path bootstrap (platform-agnostic) ──
_SCRIPT_DIR = Path(__file__).resolve().parent              # gateway/scripts/
_GATEWAY_DIR = _SCRIPT_DIR.parent                          # gateway/
_PROJECT_DIR = _GATEWAY_DIR.parent                         # AgentsMeeting/

USAGE_MONITOR_DIR = _SCRIPT_DIR / "usage_monitor"
ACCOUNTS_FILE = USAGE_MONITOR_DIR / "accounts.json"
COOKIES_DIR = USAGE_MONITOR_DIR / "cookies"
TEMP_DIR = _GATEWAY_DIR / "temp"
OUTPUT_FILE = TEMP_DIR / "usage_stats.json"
LOG_FILE = _GATEWAY_DIR / "logs" / "usage_collector.log"

# ── Ensure log dir exists before FileHandler ──
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)

# ── Logging setup ──
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [usage] %(levelname)s: %(message)s',
    handlers=[
        logging.FileHandler(str(LOG_FILE), encoding='utf-8'),
        logging.StreamHandler(sys.stdout)
    ]
)
log = logging.getLogger('usage_collector')

# ── Constants ──
DEFAULT_POLL_INTERVAL_SEC = 300   # 5 minutes
CDP_PROXY_URL = os.environ.get('CDP_PROXY_URL', 'http://localhost:3456')
CDP_EVAL_TIMEOUT = 20
CDP_TARGETS_TIMEOUT = 5
HTTP_TIMEOUT = 45  # opencode.ai 响应慢（2026-08 实测 10-12s），放宽到 45s
OPENCODE_BASE_URL = 'https://opencode.ai'

# ── SSR HTML parsing ──
#
# opencode.ai SSR hydration format:
#   lite.subscription.get["wrk_XXX"]).p.v = $R[N] = {
#     rollingUsage: $R[M] = {status:"ok", resetInSec:16602, usagePercent:37},
#     weeklyUsage:  $R[M] = {status:"ok", resetInSec:402661, usagePercent:15},
#     monthlyUsage: $R[M] = {status:"ok", resetInSec:2674309, usagePercent:7}
#   }

_METRIC_RE = re.compile(
    r'(?P<metric>rollingUsage|weeklyUsage|monthlyUsage)'
    r'\s*:\s*'
    r'(?:\$R\[\d+\]\s*=\s*)?'
    r'\{(?P<inner>[^}]*)\}',
    re.DOTALL
)

_WORKSPACE_LITERAL_RE = re.compile(r'(?P<workspace>wrk_[A-Z0-9]{20,})')


def ensure_dirs():
    """Make sure all expected directories exist."""
    USAGE_MONITOR_DIR.mkdir(parents=True, exist_ok=True)
    COOKIES_DIR.mkdir(parents=True, exist_ok=True)
    TEMP_DIR.mkdir(parents=True, exist_ok=True)
    LOG_FILE.parent.mkdir(parents=True, exist_ok=True)


def load_accounts_config():
    """Load accounts.json."""
    if not ACCOUNTS_FILE.exists():
        return None
    try:
        with open(str(ACCOUNTS_FILE), 'r', encoding='utf-8-sig') as f:
            return json.load(f)
    except Exception as e:
        log.error(f'Failed to read accounts.json: {e}')
        return None


# ═══════════════════════════════════════════════════════════════════════
#  PRIMARY MODE: HTTP + Stored Cookies
# ═══════════════════════════════════════════════════════════════════════

def load_stored_cookies(key_id):
    """
    Load cookies from cookies/{key_id}.json.
    Returns list of cookie dicts or None if file doesn't exist.
    """
    cookie_file = COOKIES_DIR / f'{key_id}.json'
    if not cookie_file.exists():
        return None
    try:
        with open(str(cookie_file), 'r', encoding='utf-8') as f:
            data = json.load(f)
            return data.get('cookies', [])
    except Exception as e:
        log.error(f'Failed to load cookies for {key_id}: {e}')
        return None


def _cookie_dict_to_jar(cookies_list, workspace_id=None):
    """Convert stored cookie dicts → http.cookiejar.CookieJar."""
    jar = http.cookiejar.CookieJar()
    for c in cookies_list:
        domain = c.get('domain', '')
        try:
            cookie = http.cookiejar.Cookie(
                version=0, name=c.get('name', ''), value=c.get('value', ''),
                port=None, port_specified=False,
                domain=domain, domain_specified=bool(domain),
                domain_initial_dot=domain.startswith('.'),
                path=c.get('path', '/'), path_specified=True,
                secure=c.get('secure', False), expires=c.get('expires'),
                discard=False, comment=None, comment_url=None,
                rest={'HttpOnly': None} if c.get('httpOnly') else {},
                rfc2109=False,
            )
            jar.set_cookie(cookie)
        except Exception as e:
            log.warning(f'Failed to convert cookie {c.get("name","?")}: {e}')
    return jar


def _jar_to_cookie_dicts(jar):
    """Convert CookieJar → list of cookie dicts for JSON storage."""
    out = []
    for c in jar:
        out.append({
            'name': c.name, 'value': c.value,
            'domain': c.domain, 'path': c.path,
            'secure': c.secure, 'expires': c.expires,
            'httpOnly': 'HttpOnly' in (c._rest or {}),
        })
    return out


def _save_cookies_back(key_id, jar, workspace_id):
    """Save updated cookies from jar back to cookies/{key_id}.json."""
    cookies_list = _jar_to_cookie_dicts(jar)
    if not cookies_list:
        return
    cookie_file = COOKIES_DIR / f'{key_id}.json'
    data = {
        'key_id': key_id, 'workspace_id': workspace_id,
        'extracted_at': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        'cookie_count': len(cookies_list), 'cookies': cookies_list,
    }
    try:
        with open(str(cookie_file), 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        log.info(f'[keepalive] Updated cookies for {key_id} ({len(cookies_list)} cookies)')
    except Exception as e:
        log.error(f'[keepalive] Failed to save cookies for {key_id}: {e}')


def fetch_ssr_via_http(workspace_id, cookies_list, key_id):
    """
    HTTP GET with CookieJar (auto-captures Set-Cookie for keepalive).
    Returns (html_text, error_message, updated_jar).
    """
    url = f'{OPENCODE_BASE_URL}/workspace/{workspace_id}/go'
    jar = _cookie_dict_to_jar(cookies_list, workspace_id)
    opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
    req = urllib.request.Request(url, method='GET')
    req.add_header('User-Agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36')
    req.add_header('Accept', 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8')

    try:
        with opener.open(req, timeout=HTTP_TIMEOUT) as resp:
            status = resp.getcode()
            final_url = resp.geturl()

            if 'auth.opencode.ai' in final_url or 'auth/authorize' in final_url:
                return None, 'session expired (redirected to auth.opencode.ai)', jar

            if status != 200:
                return None, f'HTTP {status}', jar

            html = resp.read().decode('utf-8', errors='replace')
            return html, None, jar
    except urllib.error.HTTPError as e:
        return None, f'HTTP {e.code}: {e.reason}', jar
    except urllib.error.URLError as e:
        return None, f'URL error: {e.reason}', jar
    except Exception as e:
        return None, f'Unexpected: {e}', jar


def _cookies_changed(original, jar):
    """Check if jar cookies differ from original list (Set-Cookie was received)."""
    jar_cookies = _jar_to_cookie_dicts(jar)
    if len(jar_cookies) != len(original):
        return True
    for a, b in zip(jar_cookies, original):
        if a.get('value') != b.get('value') or a.get('expires') != b.get('expires'):
            return True
    return False


def collect_via_http(key_id, account_cfg):
    """
    Collect usage data for one account via HTTP + stored cookies.
    Captures Set-Cookie for keepalive (saves refreshed cookies back to JSON).
    """
    workspace_id = account_cfg.get('workspace_id', '').strip()
    if not workspace_id:
        return {
            'key_id': key_id, 'label': account_cfg.get('label', key_id),
            'workspace_id': None, 'error': 'no workspace_id configured',
            'session_expired': False,
            'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        }

    cookies = load_stored_cookies(key_id)
    if not cookies:
        return {
            'key_id': key_id, 'label': account_cfg.get('label', key_id),
            'workspace_id': workspace_id,
            'error': 'no stored cookies — run extract_cookies.py after logging in',
            'session_expired': False,
            'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        }

    log.info(f'[HTTP] Fetching {workspace_id}/go for {key_id}...')

    html, err, jar = fetch_ssr_via_http(workspace_id, cookies, key_id)

    # Cookie keepalive: if server sent Set-Cookie, save refreshed cookies
    if _cookies_changed(cookies, jar):
        _save_cookies_back(key_id, jar, workspace_id)

    if err:
        session_expired = 'auth' in err.lower() or 'session' in err.lower()
        log.warning(f'[HTTP] {key_id} failed: {err}')
        return {
            'key_id': key_id, 'label': account_cfg.get('label', key_id),
            'workspace_id': workspace_id, 'error': err,
            'session_expired': session_expired,
            'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        }

    parsed = parse_usage_from_html(html, workspace_id)
    if not parsed:
        return {
            'key_id': key_id, 'label': account_cfg.get('label', key_id),
            'workspace_id': workspace_id,
            'error': 'no usage metrics in SSR HTML',
            'session_expired': False,
            'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        }

    # Check subscription status from page text
    sub_active = check_subscription_status(html)
    renewal_status = account_cfg.get('renewal', 'active')

    result = {
        'key_id': key_id, 'label': account_cfg.get('label', key_id),
        'workspace_id': parsed['workspace_id'],
        'rolling': parsed['rolling'], 'weekly': parsed['weekly'],
        'monthly': parsed['monthly'], 'session_expired': False,
        'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        'error': None,
        'subscribed': sub_active,
        'renewal': renewal_status,
    }

    if not sub_active:
        log.warning(f'[HTTP] {key_id} subscription NO LONGER ACTIVE — needs cleanup')
    elif renewal_status == 'cancelled':
        log.info(f'[HTTP] {key_id} renewal cancelled, but subscription still active')

    log.info(f'[HTTP] SUCCESS {key_id}: '
             f'rolling={result["rolling"]["usage_percent"]}% '
             f'weekly={result["weekly"]["usage_percent"]}% '
             f'monthly={result["monthly"]["usage_percent"]}%')

    return result


# ═══════════════════════════════════════════════════════════════════════
#  BONUS MODE: CDP (currently logged-in account only)
# ═══════════════════════════════════════════════════════════════════════

def cdp_get_targets():
    """Query CDP proxy for Chrome targets."""
    url = f'{CDP_PROXY_URL}/targets'
    try:
        req = urllib.request.Request(url, method='GET')
        with urllib.request.urlopen(req, timeout=CDP_TARGETS_TIMEOUT) as resp:
            data = json.loads(resp.read().decode('utf-8'))
            if isinstance(data, list):
                return data
            if isinstance(data, dict):
                return data.get('targets', data.get('result', []))
            return []
    except Exception as e:
        log.error(f'CDP /targets error: {e}')
        return []


def cdp_eval(target_id, js_code):
    """Execute JS in a Chrome tab via CDP proxy."""
    url = f'{CDP_PROXY_URL}/eval?target={target_id}'
    data = js_code.encode('utf-8')
    try:
        req = urllib.request.Request(url, data=data, method='POST')
        with urllib.request.urlopen(req, timeout=CDP_EVAL_TIMEOUT) as resp:
            return json.loads(resp.read().decode('utf-8'))
    except Exception as e:
        log.error(f'CDP /eval error: {e}')
        return None


def cdp_navigate(target_id, url):
    """Navigate a Chrome tab to a URL via CDP proxy."""
    full_url = f'{CDP_PROXY_URL}/navigate?target={target_id}&url={urllib.parse.quote(url, safe="")}'
    try:
        req = urllib.request.Request(full_url, method='GET')
        with urllib.request.urlopen(req, timeout=10) as resp:
            return json.loads(resp.read().decode('utf-8'))
    except Exception as e:
        log.error(f'CDP /navigate error: {e}')
        return None


def find_opencode_tab():
    """Find an opencode.ai workspace tab in Chrome."""
    targets = cdp_get_targets()
    if not targets:
        log.warning('No Chrome targets — CDP proxy may be down or Chrome not running')
        return None

    for t in targets:
        url = t.get('url', '')
        if t.get('type') == 'page' and 'opencode.ai/workspace/' in url:
            return t
    for t in targets:
        url = t.get('url', '')
        if t.get('type') == 'page' and 'opencode.ai' in url:
            return t
    for t in targets:
        if t.get('type') == 'page':
            return t
    return None


def fetch_ssr_html_via_cdp(target_id, workspace_url=None):
    """Use CDP /eval to run fetch(location.href) inside Chrome tab."""
    if workspace_url:
        log.info(f'[CDP] Navigating to {workspace_url}')
        cdp_navigate(target_id, workspace_url)
        time.sleep(3)

    js = """
(async () => {
    try {
        var r = await fetch(location.href, {credentials:'include'});
        var html = await r.text();
        var wsMatch = html.match(/wrk_[A-Z0-9]{20,}/);
        var wsId = wsMatch ? wsMatch[0] : null;
        return JSON.stringify({
            ok: true,
            status: r.status,
            url: location.href,
            html: html,
            workspace_id: wsId
        });
    } catch(e) {
        return JSON.stringify({ok: false, error: e.message});
    }
})()
""".strip()

    result = cdp_eval(target_id, js)
    if not result:
        return None, 'CDP /eval returned no result'

    value = None
    if isinstance(result, dict):
        value = result.get('value') or result.get('result')
    elif isinstance(result, str):
        value = result

    if not value:
        return None, f'CDP /eval returned empty: {result}'

    try:
        parsed = json.loads(value)
    except (json.JSONDecodeError, TypeError) as e:
        return None, f'Failed to parse /eval return: {e}'

    if not parsed.get('ok'):
        return None, f'JS fetch failed: {parsed.get("error", "unknown")}'

    html = parsed.get('html', '')
    ws_id = parsed.get('workspace_id')
    status = parsed.get('status', 0)
    url = parsed.get('url', '')

    log.info(f'[CDP] fetch: status={status} url={url[:120]} html_len={len(html)} ws_id={ws_id}')

    if 'auth.opencode.ai' in url or 'auth/authorize' in url:
        return None, 'session expired (CDP tab on auth page)'
    if status != 200 or not html:
        return None, f'CDP fetch status={status} html_len={len(html)}'
    return ws_id, html


def collect_via_cdp(accounts_cfg):
    """
    Collect usage for the currently logged-in Chrome account via CDP.
    Returns (key_id, result_dict) or (None, error_dict).
    """
    tab = find_opencode_tab()
    if not tab:
        return None, {
            'error': 'no Chrome tab (CDP proxy down or Chrome not running)',
            'session_expired': False,
        }

    target_id = tab.get('targetId') or tab.get('id') or ''
    if not target_id:
        return None, {'error': 'tab has no targetId', 'session_expired': False}

    tab_url = tab.get('url', '')

    # If on auth page, try navigating to a workspace URL
    workspace_url = None
    if 'auth.opencode.ai' in tab_url or 'auth/authorize' in tab_url:
        if accounts_cfg:
            for acc in accounts_cfg.get('accounts', []):
                ws_id = acc.get('workspace_id', '').strip()
                if ws_id:
                    workspace_url = f'{OPENCODE_BASE_URL}/workspace/{ws_id}/go'
                    break
        if not workspace_url:
            return None, {'error': 'session expired + no workspace URLs', 'session_expired': True}
    elif 'opencode.ai/workspace/' in tab_url:
        workspace_url = None  # fetch(location.href) will work

    ws_id, err_or_html = fetch_ssr_html_via_cdp(target_id, workspace_url)
    if not ws_id:
        error_msg = err_or_html or 'unknown fetch error'
        return None, {
            'error': error_msg,
            'session_expired': 'session' in error_msg.lower() or 'auth' in error_msg.lower(),
        }

    html_text = err_or_html
    parsed = parse_usage_from_html(html_text, ws_id)
    if not parsed:
        return ws_id, {
            'workspace_id': ws_id,
            'error': 'no usage metrics (no Go subscription?)',
            'session_expired': False,
        }

    # Match to account
    label = ws_id
    key_id = ws_id
    if accounts_cfg:
        for acc in accounts_cfg.get('accounts', []):
            if acc.get('workspace_id') == ws_id:
                label = acc.get('label', ws_id)
                key_id = acc.get('key_id', ws_id)
                break

    # Check subscription status (CDP fetches via fetch(), not document.body, so check html_text)
    sub_active = check_subscription_status(html_text)
    renewal_status = 'active'
    if accounts_cfg:
        for acc in accounts_cfg.get('accounts', []):
            if acc.get('workspace_id') == ws_id:
                renewal_status = acc.get('renewal', 'active')
                break

    result = {
        'key_id': key_id,
        'label': label,
        'workspace_id': parsed['workspace_id'],
        'rolling': parsed['rolling'],
        'weekly': parsed['weekly'],
        'monthly': parsed['monthly'],
        'session_expired': False,
        'last_update_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        'error': None,
        'subscribed': sub_active,
        'renewal': renewal_status,
    }

    if not sub_active:
        log.warning(f'[CDP] {key_id} subscription NO LONGER ACTIVE — needs cleanup')
    elif renewal_status == 'cancelled':
        log.info(f'[CDP] {key_id} renewal cancelled, but subscription still active')

    log.info(f'[CDP] SUCCESS {key_id}: '
             f'rolling={result["rolling"]["usage_percent"]}% '
             f'weekly={result["weekly"]["usage_percent"]}% '
             f'monthly={result["monthly"]["usage_percent"]}%')

    return key_id, result


# ═══════════════════════════════════════════════════════════════════════
#  SSR Parsing (shared by both modes)
# ═══════════════════════════════════════════════════════════════════════

def parse_usage_from_html(html_text, workspace_id_hint):
    """
    Parse SSR HTML for usage data using per-metric regex.
    Returns dict with workspace_id, rolling, weekly, monthly or None.
    """
    metrics = {}
    for m in _METRIC_RE.finditer(html_text):
        metric_name = m.group('metric')
        if metric_name in metrics:
            continue
        inner = m.group('inner')
        status_match = re.search(r'status\s*:\s*"([^"]*)"', inner)
        pct_match = re.search(r'usagePercent\s*:\s*(\d+)', inner)
        reset_match = re.search(r'resetInSec\s*:\s*(\d+)', inner)
        metrics[metric_name] = {
            'usage_percent': int(pct_match.group(1)) if pct_match else None,
            'reset_in_sec': int(reset_match.group(1)) if reset_match else None,
            'status': status_match.group(1) if status_match else None,
        }

    ws_match = _WORKSPACE_LITERAL_RE.search(html_text)
    workspace_id = ws_match.group('workspace') if ws_match else workspace_id_hint

    if not metrics:
        return None

    log.info(f'Parsed {workspace_id}: '
             f'rolling={metrics.get("rollingUsage", {}).get("usage_percent")}% '
             f'weekly={metrics.get("weeklyUsage", {}).get("usage_percent")}% '
             f'monthly={metrics.get("monthlyUsage", {}).get("usage_percent")}%')

    return {
        'workspace_id': workspace_id,
        'rolling': metrics.get('rollingUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
        'weekly': metrics.get('weeklyUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
        'monthly': metrics.get('monthlyUsage', {'usage_percent': None, 'reset_in_sec': None, 'status': None}),
    }


def check_subscription_status(html_text):
    """
    Check whether the SSR HTML page indicates an active subscription.

    Looks for 'subscribed' text (case-insensitive) in the HTML.
    Returns True if the user is still subscribed, False otherwise.
    """
    # Heuristic: page is loaded (no auth redirect) + contains "subscribed" keyword
    if not html_text:
        return False
    lower = html_text.lower()
    # Check for positive signal
    subscribed = 'subscribed' in lower
    return subscribed


# ═══════════════════════════════════════════════════════════════════════
#  Cache merge logic
# ═══════════════════════════════════════════════════════════════════════

def load_existing_cache():
    """Load existing usage_stats.json or return empty structure."""
    if not OUTPUT_FILE.exists():
        return {'accounts': {}}
    try:
        with open(str(OUTPUT_FILE), 'r', encoding='utf-8') as f:
            data = json.load(f)
            if 'accounts' in data and isinstance(data['accounts'], list):
                accounts_dict = {}
                for acc in data['accounts']:
                    kid = acc.get('key_id', acc.get('workspace_id', '?'))
                    accounts_dict[kid] = acc
                data['accounts'] = accounts_dict
            return data
    except Exception:
        return {'accounts': {}}


def merge_result_into_cache(cache, key_id, new_result):
    """Update cache with new result for one account."""
    if 'accounts' not in cache or not isinstance(cache['accounts'], dict):
        cache['accounts'] = {}
    cache['accounts'][key_id] = new_result


# ═══════════════════════════════════════════════════════════════════════
#  Main collection (Hybrid)
# ═══════════════════════════════════════════════════════════════════════

def collect_all(cdp_only=False):
    """
    Hybrid collection: HTTP for all accounts + CDP bonus for current login.
    """
    ensure_dirs()
    accounts_cfg = load_accounts_config()
    if not accounts_cfg:
        log.error('No accounts.json found')
        return {'ok': False, 'error': 'no accounts.json', 'accounts': []}

    cache = load_existing_cache()
    collected_count = 0
    cdp_key_id = None

    # ── PRIMARY: HTTP collection for OpenCode Go accounts only ──
    if not cdp_only:
        for acc_cfg in accounts_cfg.get('accounts', []):
            key_id = acc_cfg.get('key_id', '')
            provider = acc_cfg.get('provider', 'opencode_go')
            if not key_id:
                continue
            # Skip non-OCG providers (Kimi etc.)
            if provider != 'opencode_go':
                continue
            result = collect_via_http(key_id, acc_cfg)
            merge_result_into_cache(cache, key_id, result)

            if result.get('error') is None:
                collected_count += 1
            elif result.get('session_expired'):
                log.warning(f'{key_id} session expired — needs re-login + re-extract cookies')

    # ── BONUS: CDP collection for currently logged-in account ──
    try:
        cdp_kid, cdp_result = collect_via_cdp(accounts_cfg)
        if cdp_kid:
            cdp_key_id = cdp_kid
            # CDP result overrides HTTP for this account (more reliable — live cookies)
            merge_result_into_cache(cache, cdp_kid, cdp_result)
            if cdp_result.get('error') is None:
                collected_count += 1
                log.info(f'CDP override for {cdp_kid}')
    except Exception as e:
        log.warning(f'CDP bonus collection skipped: {e}')

    # ── Build output (OpenCode Go accounts only) ──
    accounts_list = []
    for acc_cfg in accounts_cfg.get('accounts', []):
        # Skip non-OCG providers
        if acc_cfg.get('provider', 'opencode_go') != 'opencode_go':
            continue
        kid = acc_cfg.get('key_id', acc_cfg.get('workspace_id', '?'))
        if kid in cache.get('accounts', {}):
            # Merge in account config fields
            cached = cache['accounts'][kid]
            cached['key_id'] = kid
            cached['label'] = acc_cfg.get('label', kid)
            cached['provider'] = acc_cfg.get('provider', 'opencode_go')
            cached['workspace_id'] = acc_cfg.get('workspace_id') or cached.get('workspace_id')
            accounts_list.append(cached)
        else:
            accounts_list.append({
                'key_id': kid,
                'label': acc_cfg.get('label', kid),
                'workspace_id': acc_cfg.get('workspace_id'),
                'rolling': None,
                'weekly': None,
                'monthly': None,
                'session_expired': False,
                'last_update_iso': None,
                'error': 'not yet collected',
            })

    output = {
        'ok': collected_count > 0,
        'last_refresh_iso': datetime.now(timezone.utc).isoformat(timespec='seconds'),
        'account_count': len(accounts_list),
        'collected_count': collected_count,
        'cdp_override_key_id': cdp_key_id,
        'accounts': accounts_list,
    }

    save_result(output)

    ok = sum(1 for a in accounts_list if a.get('rolling') and a['rolling'].get('usage_percent') is not None)
    err = sum(1 for a in accounts_list if a.get('error'))
    log.info(f'Done: {len(accounts_list)} accounts, {ok} with data, {err} errored.')

    return output


def save_result(data):
    """Persist to gateway/temp/usage_stats.json."""
    ensure_dirs()
    try:
        with open(str(OUTPUT_FILE), 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        log.info(f'Saved to {OUTPUT_FILE}')
    except Exception as e:
        log.error(f'Failed to save: {e}')


def load_result():
    """Return cached result or empty dict."""
    if not OUTPUT_FILE.exists():
        return {
            'ok': False,
            'error': 'no usage_stats.json yet',
            'last_refresh_iso': None,
            'accounts': []
        }
    try:
        with open(str(OUTPUT_FILE), 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        return {'ok': False, 'error': str(e), 'accounts': []}


def run_once(cdp_only=False):
    """One-shot collection."""
    return collect_all(cdp_only=cdp_only)


def run_daemon(interval=DEFAULT_POLL_INTERVAL_SEC):
    """Daemon mode."""
    ensure_dirs()
    log.info(f'Daemon started, interval={interval}s')
    while True:
        try:
            collect_all()
        except Exception as e:
            log.exception(f'Daemon error: {e}')
        time.sleep(interval)


def main():
    parser = argparse.ArgumentParser(description='OpenCode Go usage collector (Hybrid v3.0)')
    parser.add_argument('--daemon', action='store_true', help='Run as daemon')
    parser.add_argument('--interval', type=int, default=DEFAULT_POLL_INTERVAL_SEC,
                        help=f'Polling interval in seconds (default: {DEFAULT_POLL_INTERVAL_SEC})')
    parser.add_argument('--print', action='store_true', help='Print result JSON')
    parser.add_argument('--cdp-only', action='store_true', help='Only collect via CDP (skip HTTP)')
    args = parser.parse_args()

    if args.daemon:
        run_daemon(interval=args.interval)
    else:
        result = run_once(cdp_only=args.cdp_only)
        if args.print:
            print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()