#!/usr/bin/env python3
"""
extract_cookies.py — 提取并保存 opencode.ai cookies (含 httpOnly) 到 per-account JSON
===============================================================================
通过 CDP proxy /cookies 端点提取 Chrome 当前登录账号的全部 cookies，
过滤 opencode.ai domain (含 auth.opencode.ai)，保存到 cookies/{key_id}.json。

用法:
  python extract_cookies.py <key_id> [--workspace-id <wrk_XXX>]
  python extract_cookies.py key1 --workspace-id wrk_01KQT521KYE2P2QRZ10N8MVPXA

依赖: 仅 Python stdlib (urllib, json) + CDP proxy localhost:3456
"""
import sys
import json
import urllib.request
from pathlib import Path
from datetime import datetime, timezone

CDP_PROXY_URL = 'http://localhost:3456'
COOKIES_DIR = Path(__file__).resolve().parent / 'cookies'
ACCOUNTS_FILE = Path(__file__).resolve().parent / 'accounts.json'


def fetch_cookies_from_cdp():
    """Call CDP proxy /cookies endpoint, return list of cookie dicts."""
    url = f'{CDP_PROXY_URL}/cookies'
    req = urllib.request.Request(url, method='GET')
    with urllib.request.urlopen(req, timeout=15) as resp:
        data = json.loads(resp.read().decode('utf-8'))
        if isinstance(data, list):
            return data
        return data.get('cookies', data.get('result', {}).get('cookies', []))


def filter_opencode_cookies(cookies):
    """Filter to opencode.ai + auth.opencode.ai domain cookies."""
    return [c for c in cookies if c.get('domain', '').endswith('opencode.ai')]


def save_cookies(key_id, cookies, workspace_id=None):
    """Save cookies to cookies/{key_id}.json along with metadata."""
    COOKIES_DIR.mkdir(parents=True, exist_ok=True)
    out_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),
        'cookies': cookies,
    }
    with open(str(out_file), 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    return out_file


def update_accounts_json(key_id, workspace_id):
    """Update the workspace_id for the given key_id in accounts.json."""
    if not ACCOUNTS_FILE.exists() or not workspace_id:
        return
    try:
        with open(str(ACCOUNTS_FILE), 'r', encoding='utf-8-sig') as f:
            cfg = json.load(f)
    except Exception:
        return
    updated = False
    for acc in cfg.get('accounts', []):
        if acc.get('key_id') == key_id:
            if not acc.get('workspace_id') and workspace_id:
                acc['workspace_id'] = workspace_id
                updated = True
            break
    if updated:
        with open(str(ACCOUNTS_FILE), 'w', encoding='utf-8') as f:
            json.dump(cfg, f, ensure_ascii=False, indent=2)
        print(f'Updated accounts.json: {key_id} → workspace_id={workspace_id}')


def main():
    if len(sys.argv) < 2:
        print('Usage: python extract_cookies.py <key_id> [--workspace-id <wrk_XXX>]')
        sys.exit(1)

    key_id = sys.argv[1]
    workspace_id = None
    if '--workspace-id' in sys.argv:
        idx = sys.argv.index('--workspace-id')
        workspace_id = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else None

    print(f'Extracting cookies for {key_id}...')
    try:
        all_cookies = fetch_cookies_from_cdp()
    except Exception as e:
        print(f'ERROR: CDP /cookies failed: {e}')
        sys.exit(1)

    print(f'Total cookies from Chrome: {len(all_cookies)}')
    oc_cookies = filter_opencode_cookies(all_cookies)
    print(f'opencode.ai cookies: {len(oc_cookies)}')
    for c in oc_cookies:
        print(f'  {c["name"]}: domain={c["domain"]} httpOnly={c.get("httpOnly", False)} len={len(c.get("value", ""))}')

    if not oc_cookies:
        print('ERROR: No opencode.ai cookies found — is Chrome logged in to an opencode account?')
        sys.exit(1)

    out_file = save_cookies(key_id, oc_cookies, workspace_id)
    print(f'Saved {len(oc_cookies)} cookies to {out_file}')

    if workspace_id:
        update_accounts_json(key_id, workspace_id)
        print(f'Workspace ID: {workspace_id}')


if __name__ == '__main__':
    main()