#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""批量重铸 CPA auth：补回 grok-cli:access scope。

只覆盖写回，不删除原文件。默认仅处理缺 grok-cli:access 的 xai-*.json。

用法：
  python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --dry-run
  python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --limit 5
  python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --concurrency 6
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

_ROOT = Path(__file__).resolve().parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))


def _jwt_payload(token: str) -> dict[str, Any]:
    parts = (token or "").split(".")
    if len(parts) < 2:
        raise ValueError("invalid JWT")
    seg = parts[1] + "=" * (-len(parts[1]) % 4)
    return json.loads(base64.urlsafe_b64decode(seg.encode("ascii")))


def _scope_set(token: str) -> set[str]:
    try:
        return set(str(_jwt_payload(token).get("scope") or "").split())
    except Exception:
        return set()


def _needs_remint(data: dict[str, Any], *, only_missing_cli: bool) -> tuple[bool, str]:
    at = str(data.get("access_token") or "").strip()
    if not at:
        return False, "no_access_token"
    scopes = _scope_set(at)
    has_cli = "grok-cli:access" in scopes
    if only_missing_cli and has_cli:
        return False, "already_has_cli_scope"
    sso = str(data.get("sso") or "").strip()
    if not sso:
        return False, "no_sso"
    return True, "missing_cli_scope" if not has_cli else "force"


def _atomic_write(path: Path, payload: dict[str, Any]) -> None:
    tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}.{threading.get_ident()}")
    text = json.dumps(payload, ensure_ascii=False, indent=2) + "\n"
    tmp.write_text(text, encoding="utf-8")
    os.replace(tmp, path)


def remint_one(
    path: Path,
    *,
    proxy: str | None,
    dry_run: bool,
    only_missing_cli: bool,
) -> dict[str, Any]:
    import cpa
    from oidc_mint.oauth_code import GROK_CLI_SCOPE, mint_from_sso, normalize_sso_cookie

    try:
        data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "path": str(path), "error": f"read: {exc}"}

    need, reason = _needs_remint(data, only_missing_cli=only_missing_cli)
    email = str(data.get("email") or "").strip()
    if not need:
        return {
            "ok": True,
            "skipped": True,
            "path": str(path),
            "email": email,
            "reason": reason,
        }

    sso = normalize_sso_cookie(str(data.get("sso") or ""))
    if dry_run:
        return {
            "ok": True,
            "dry_run": True,
            "path": str(path),
            "email": email,
            "reason": reason,
        }

    try:
        tokens = mint_from_sso(
            sso,
            proxy=proxy,
            require_referrer=True,
            log=lambda m: None,
        )
    except Exception as exc:  # noqa: BLE001
        return {
            "ok": False,
            "path": str(path),
            "email": email,
            "error": f"mint: {exc}",
            "reason": reason,
        }

    access = str(tokens.get("access_token") or "").strip()
    refresh = str(tokens.get("refresh_token") or "").strip()
    if GROK_CLI_SCOPE not in _scope_set(access):
        return {
            "ok": False,
            "path": str(path),
            "email": email,
            "error": f"minted token still missing {GROK_CLI_SCOPE}",
            "reason": reason,
        }

    try:
        payload = cpa.build_cpa_xai_auth(
            email=email or str(data.get("email") or ""),
            access_token=access,
            refresh_token=refresh,
            id_token=tokens.get("id_token") or data.get("id_token"),
            expires_in=tokens.get("expires_in"),
            base_url=data.get("base_url") or cpa.CLI_BASE_URL,
            sso=sso,
            headers=data.get("headers") if isinstance(data.get("headers"), dict) else None,
        )
        # 保留原 disabled 状态，避免重铸把人工禁用号重新启用
        if "disabled" in data:
            payload["disabled"] = bool(data.get("disabled"))
        _atomic_write(path, payload)
    except Exception as exc:  # noqa: BLE001
        return {
            "ok": False,
            "path": str(path),
            "email": email,
            "error": f"write: {exc}",
            "reason": reason,
        }

    return {
        "ok": True,
        "path": str(path),
        "email": email or payload.get("email"),
        "reason": reason,
        "referrer": payload.get("referrer"),
        "scope_ok": True,
    }


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Remint CPA auths missing grok-cli:access")
    p.add_argument(
        "--auth-dir",
        default=os.environ.get("CPA_AUTH_DIR", str(_ROOT / "cpa_auths")),
        help="CPA auth 目录（默认 ./cpa_auths 或 $CPA_AUTH_DIR）",
    )
    p.add_argument("--proxy", default=os.environ.get("https_proxy") or os.environ.get("http_proxy") or "")
    p.add_argument("--concurrency", type=int, default=4)
    p.add_argument("--limit", type=int, default=0, help="最多处理 N 个需重铸文件（0=不限）")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument(
        "--all",
        action="store_true",
        help="强制全部重铸（默认只处理缺 grok-cli:access 的）",
    )
    p.add_argument(
        "--log-file",
        default="",
        help="结果日志路径（默认 auth-dir/remint_cli_scope.log）",
    )
    args = p.parse_args(argv)

    auth_dir = Path(args.auth_dir).expanduser().resolve()
    if not auth_dir.is_dir():
        print(f"[!] auth-dir 不存在: {auth_dir}", flush=True)
        return 2

    proxy = (args.proxy or "").strip() or None
    only_missing_cli = not args.all
    files = sorted(auth_dir.glob("xai-*.json"))
    targets: list[Path] = []
    skipped = 0
    for path in files:
        try:
            data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
            need, _ = _needs_remint(data, only_missing_cli=only_missing_cli)
            if need:
                targets.append(path)
            else:
                skipped += 1
        except Exception:
            skipped += 1
        if args.limit and len(targets) >= args.limit:
            break

    log_path = Path(args.log_file).expanduser() if args.log_file else (auth_dir / "remint_cli_scope.log")
    print(
        f"[*] auth-dir={auth_dir} total={len(files)} targets={len(targets)} "
        f"skipped_scan={skipped} concurrency={args.concurrency} "
        f"proxy={proxy or '(direct)'} dry_run={args.dry_run}",
        flush=True,
    )
    if not targets:
        print("[*] 无需重铸", flush=True)
        return 0

    ok_n = 0
    fail_n = 0
    skip_n = 0
    lock = threading.Lock()
    started = time.time()

    def _handle(path: Path) -> dict[str, Any]:
        return remint_one(
            path,
            proxy=proxy,
            dry_run=args.dry_run,
            only_missing_cli=only_missing_cli,
        )

    def _record(res: dict[str, Any]) -> None:
        nonlocal ok_n, fail_n, skip_n
        line = json.dumps(res, ensure_ascii=False)
        with lock:
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(line + "\n")
            if res.get("skipped") or res.get("dry_run"):
                skip_n += 1
                tag = "SKIP" if res.get("skipped") else "DRY"
            elif res.get("ok"):
                ok_n += 1
                tag = "OK"
            else:
                fail_n += 1
                tag = "FAIL"
            email = res.get("email") or Path(str(res.get("path") or "")).name
            extra = res.get("error") or res.get("reason") or ""
            print(f"[{tag}] {email} {extra}", flush=True)

    # 清空/追加日志头
    with open(log_path, "a", encoding="utf-8") as f:
        f.write(
            json.dumps(
                {
                    "event": "start",
                    "ts": int(time.time()),
                    "targets": len(targets),
                    "dry_run": bool(args.dry_run),
                    "proxy": proxy or "",
                },
                ensure_ascii=False,
            )
            + "\n"
        )

    if args.concurrency <= 1:
        for path in targets:
            _record(_handle(path))
    else:
        with ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as ex:
            futs = [ex.submit(_handle, path) for path in targets]
            for fut in as_completed(futs):
                try:
                    _record(fut.result())
                except Exception as exc:  # noqa: BLE001
                    _record({"ok": False, "error": f"worker: {exc}"})

    elapsed = time.time() - started
    print(
        f"[*] done ok={ok_n} fail={fail_n} skip/dry={skip_n} "
        f"elapsed={elapsed:.1f}s log={log_path}",
        flush=True,
    )
    return 0 if fail_n == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())