"""xAI OAuth Authorization Code + PKCE (SSO cookie → CPA token).

对齐最新可用流程：authorize / consent 必须带 referrer=grok-build，
否则 access_token JWT 缺少 referrer 字段，cli-chat-proxy / grok-build 不可用。

参考实现：sso -> oauth2/authorize(referrer=grok-build) -> consent allow
        -> oauth2/token (authorization_code + PKCE)
"""

from __future__ import annotations

import base64
import hashlib
import json
import re
import secrets
import time
from dataclasses import dataclass
from typing import Any, Callable
from urllib.parse import parse_qs, urlencode, urljoin, urlparse

from .proxyutil import resolve_proxy

CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"
ISSUER = "https://auth.x.ai"
TOKEN_URL = f"{ISSUER}/oauth2/token"
AUTHORIZE_URL = f"{ISSUER}/oauth2/authorize"
REDIRECT_URI = "http://127.0.0.1:56121/callback"
# 比旧 device-code scope 多 conversations:*，对齐 grok-build
SCOPE = (
    "openid profile email offline_access "
    "grok-cli:access api:access conversations:read conversations:write"
)
GROK_REFERRER = "grok-build"
GROK_VERSION = "0.2.93"
GROK_TOKEN_UA = (
    f"grok-pager/{GROK_VERSION} grok-shell/{GROK_VERSION} (linux; x86_64)"
)
# Next.js Server Action id（consent 页 POST 需要）
NEXT_ACTION_ID = "4005315a1d7e426de592990bb54bb37471f39dd6d2"
BROWSER_UA = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
)

LogFn = Callable[[str], None]


def _noop_log(_: str) -> None:
    return None


class OAuthCodeError(RuntimeError):
    pass


@dataclass
class AuthCodeFlow:
    state: str
    nonce: str
    code_verifier: str
    code_challenge: str


@dataclass
class TokenResult:
    access_token: str
    refresh_token: str
    id_token: str | None
    token_type: str
    expires_in: int
    raw: dict[str, Any]
    referrer: str = ""


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def new_auth_code_flow() -> AuthCodeFlow:
    verifier = _b64url(secrets.token_bytes(32))
    state = _b64url(secrets.token_bytes(16))
    nonce = _b64url(secrets.token_bytes(16))
    challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
    return AuthCodeFlow(
        state=state,
        nonce=nonce,
        code_verifier=verifier,
        code_challenge=challenge,
    )


def normalize_sso_cookie(raw: str) -> str:
    token = (raw or "").strip()
    if token.lower().startswith("sso="):
        token = token[4:].strip()
    return token


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


def _short(value: str, limit: int = 240) -> str:
    value = value or ""
    return value if len(value) <= limit else value[:limit]


def _is_curl_tls_broken(exc: BaseException | str) -> bool:
    """识别 curl_cffi / libcurl OpenSSL 损坏类错误（常见于部分 Windows 环境）。"""
    text = str(exc or "").lower()
    needles = (
        "openssl_internal:invalid library",
        "tls connect error",
        "curl: (35)",
        "failed to perform, curl: (35)",
        "ssl_error_syscall",
        "ssl connect error",
        "wrong version number",
    )
    return any(n in text for n in needles)


def _make_std_session(proxy: str | None = None):
    try:
        import requests as std_requests
    except ImportError as e:  # pragma: no cover
        raise OAuthCodeError(
            "需要 curl_cffi 或 requests 才能执行 SSO→OAuth 转换"
        ) from e
    resolved = resolve_proxy(proxy)
    proxies = {"http": resolved, "https": resolved} if resolved else None
    sess = std_requests.Session()
    if proxies:
        sess.proxies.update(proxies)
    try:
        sess._cpa_http_backend = "requests"  # type: ignore[attr-defined]
    except Exception:
        pass
    return sess


def _make_curl_session(proxy: str | None = None):
    from curl_cffi import requests as crequests

    resolved = resolve_proxy(proxy)
    proxies = {"http": resolved, "https": resolved} if resolved else None
    last_err: Exception | None = None
    for impersonate in ("chrome131", "chrome124", "chrome120", "chrome110", None):
        try:
            if impersonate:
                sess = crequests.Session(impersonate=impersonate, proxies=proxies)
            else:
                sess = crequests.Session(proxies=proxies)
            try:
                sess._cpa_http_backend = f"curl_cffi:{impersonate or 'default'}"  # type: ignore[attr-defined]
            except Exception:
                pass
            return sess
        except Exception as exc:  # noqa: BLE001
            last_err = exc
            continue
    if last_err:
        raise last_err
    raise OAuthCodeError("curl_cffi Session 创建失败")


def _make_session(proxy: str | None = None, *, prefer: str = "curl"):
    """优先 curl_cffi（Chrome TLS）；prefer=requests 时直接标准库。"""
    prefer = (prefer or "curl").strip().lower()
    errors: list[str] = []
    if prefer != "requests":
        try:
            return _make_curl_session(proxy)
        except ImportError:
            errors.append("curl_cffi 未安装")
        except Exception as exc:  # noqa: BLE001
            errors.append(f"curl_cffi: {exc}")
    try:
        return _make_std_session(proxy)
    except Exception as exc:  # noqa: BLE001
        errors.append(f"requests: {exc}")
        raise OAuthCodeError(
            "无法创建 HTTP Session: " + " | ".join(errors)
        ) from exc


def _set_sso_cookies(session: Any, sso: str) -> None:
    sso = normalize_sso_cookie(sso)
    if not sso:
        raise OAuthCodeError("sso cookie 为空")
    # curl_cffi / requests cookie jar
    for domain in ("accounts.x.ai", "auth.x.ai", ".x.ai"):
        for name in ("sso", "sso-rw"):
            try:
                session.cookies.set(name, sso, domain=domain, path="/")
            except Exception:
                try:
                    session.cookies.set(name, sso)
                except Exception:
                    pass


def _browser_headers(method: str, url: str, next_action: str = "") -> dict[str, str]:
    headers = {
        "User-Agent": BROWSER_UA,
        "Sec-CH-UA": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"',
        "Sec-CH-UA-Mobile": "?0",
        "Sec-CH-UA-Platform": '"Linux"',
        "Accept-Language": "en-US,en;q=0.9",
    }
    if method.upper() == "POST":
        headers.update(
            {
                "Accept": "text/x-component",
                "Content-Type": "text/plain;charset=UTF-8",
                "Origin": "https://accounts.x.ai",
                "Referer": url,
                "Sec-Fetch-Site": "same-origin",
                "Sec-Fetch-Mode": "cors",
                "Sec-Fetch-Dest": "empty",
            }
        )
        if next_action:
            headers["Next-Action"] = next_action
    else:
        headers.update(
            {
                "Accept": (
                    "text/html,application/xhtml+xml,application/xml;q=0.9,"
                    "application/json;q=0.8,*/*;q=0.7"
                ),
                "Upgrade-Insecure-Requests": "1",
                "Sec-Fetch-Site": "none",
                "Sec-Fetch-Mode": "navigate",
                "Sec-Fetch-Dest": "document",
            }
        )
    return headers


def _token_headers() -> dict[str, str]:
    return {
        "User-Agent": GROK_TOKEN_UA,
        "Accept": "*/*",
        "X-Grok-Client-Version": GROK_VERSION,
        "Content-Type": "application/x-www-form-urlencoded",
    }


def _final_url(resp: Any) -> str:
    try:
        return str(getattr(resp, "url", "") or "")
    except Exception:
        return ""


# HTTP 铸造单步超时：直连/跨境链路差时不要卡 30s，尽快让上层走浏览器
HTTP_STEP_TIMEOUT = 12


def open_authorize_page(session: Any, flow: AuthCodeFlow) -> str:
    params = {
        "response_type": "code",
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "scope": SCOPE,
        "code_challenge": flow.code_challenge,
        "code_challenge_method": "S256",
        "state": flow.state,
        "nonce": flow.nonce,
        "referrer": GROK_REFERRER,
    }
    url = f"{AUTHORIZE_URL}?{urlencode(params)}"
    resp = session.get(
        url,
        headers=_browser_headers("GET", url),
        allow_redirects=True,
        timeout=HTTP_STEP_TIMEOUT,
    )
    body = resp.text or ""
    final = _final_url(resp)
    if resp.status_code < 200 or resp.status_code >= 300:
        raise OAuthCodeError(
            f"authorize HTTP {resp.status_code}: {_short(body)}"
        )
    if "sign-in" in final or "sign-up" in final:
        raise OAuthCodeError("sso 无效（authorize 跳转登录页）")
    if "/oauth2/consent" not in final:
        # 少数情况 consent 在 body 的 redirect 里
        m = re.search(r'https?://[^"\']+/oauth2/consent[^"\']*', body)
        if m:
            final = m.group(0)
        else:
            raise OAuthCodeError(f"authorize 未进入 consent: {final or _short(body)}")
    return final


def parse_consent_code(body: str) -> str:
    """从 Next.js RSC / text-x-component 响应中解析 code。"""
    text = body or ""
    # 1) 逐行 JSON（Go 实现路径）
    for line in text.splitlines():
        idx = line.find("{")
        if idx < 0:
            continue
        try:
            obj = json.loads(line[idx:])
        except Exception:
            continue
        if isinstance(obj, dict) and obj.get("code"):
            if obj.get("success") is False:
                raise OAuthCodeError(
                    f"consent 失败: {obj.get('error') or obj.get('action')}"
                )
            return str(obj["code"]).strip()
        if isinstance(obj, list):
            for item in obj:
                if isinstance(item, dict) and item.get("code"):
                    return str(item["code"]).strip()

    # 2) 宽松正则
    m = re.search(r'"code"\s*:\s*"([A-Za-z0-9._~\-]+)"', text)
    if m:
        return m.group(1)

    # 3) redirect 里带 code=
    m = re.search(r"[?&]code=([A-Za-z0-9._~\-]+)", text)
    if m:
        return m.group(1)

    raise OAuthCodeError(f"consent 响应缺少 code: {_short(text, 300)}")


def approve_authorization(session: Any, consent_url: str, flow: AuthCodeFlow) -> str:
    payload = [
        {
            "action": "allow",
            "clientId": CLIENT_ID,
            "redirectUri": REDIRECT_URI,
            "scope": SCOPE,
            "state": flow.state,
            "codeChallenge": flow.code_challenge,
            "codeChallengeMethod": "S256",
            "nonce": flow.nonce,
            "principalType": "User",
            "principalId": "",
            "referrer": GROK_REFERRER,
        }
    ]
    body = json.dumps(payload, separators=(",", ":"))
    resp = session.post(
        consent_url,
        data=body.encode("utf-8"),
        headers=_browser_headers("POST", consent_url, NEXT_ACTION_ID),
        allow_redirects=True,
        timeout=HTTP_STEP_TIMEOUT,
    )
    text = resp.text or ""
    if resp.status_code < 200 or resp.status_code >= 300:
        raise OAuthCodeError(f"consent HTTP {resp.status_code}: {_short(text, 300)}")

    # 有时 302 到 redirect_uri?code=
    final = _final_url(resp)
    if "code=" in final:
        qs = parse_qs(urlparse(final).query)
        code = (qs.get("code") or [""])[0]
        if code:
            return code

    return parse_consent_code(text)


def exchange_auth_code(session: Any, code: str, flow: AuthCodeFlow) -> TokenResult:
    form = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": REDIRECT_URI,
        "client_id": CLIENT_ID,
        "code_verifier": flow.code_verifier,
    }
    resp = session.post(
        TOKEN_URL,
        data=urlencode(form),
        headers=_token_headers(),
        timeout=HTTP_STEP_TIMEOUT,
    )
    text = resp.text or ""
    if resp.status_code < 200 or resp.status_code >= 300:
        raise OAuthCodeError(f"token HTTP {resp.status_code}: {_short(text, 300)}")
    try:
        data = resp.json()
    except Exception as e:
        raise OAuthCodeError(f"token 响应非 JSON: {_short(text)}") from e
    if not isinstance(data, dict) or not data.get("access_token"):
        raise OAuthCodeError(f"token 响应缺少 access_token: {data!r}")

    access = str(data["access_token"]).strip()
    refresh = str(data.get("refresh_token") or "").strip()
    if not refresh:
        raise OAuthCodeError("token 响应缺少 refresh_token")

    referrer = ""
    try:
        pl = jwt_payload(access)
        referrer = str(pl.get("referrer") or "")
    except Exception:
        pl = {}

    expires_in = int(data.get("expires_in") or 21600)
    return TokenResult(
        access_token=access,
        refresh_token=refresh,
        id_token=(str(data["id_token"]).strip() if data.get("id_token") else None),
        token_type=str(data.get("token_type") or "Bearer"),
        expires_in=expires_in,
        raw=data,
        referrer=referrer,
    )


def _run_sso_flow(
    sso: str,
    *,
    session: Any,
    log: LogFn,
    require_referrer: bool,
) -> TokenResult:
    flow = new_auth_code_flow()
    backend = getattr(session, "_cpa_http_backend", "unknown")
    log(f"Authorization Code Flow referrer={GROK_REFERRER} http={backend}")
    _set_sso_cookies(session, sso)
    consent_url = open_authorize_page(session, flow)
    log(f"authorize -> consent: {_short(consent_url, 120)}")
    code = approve_authorization(session, consent_url, flow)
    log("consent allow ok")
    token = exchange_auth_code(session, code, flow)
    if token.referrer != GROK_REFERRER:
        msg = f"access_token 未包含预期 referrer（got={token.referrer!r}）"
        if require_referrer:
            raise OAuthCodeError(msg)
        log(f"WARN {msg}")
    else:
        log("access_token referrer=grok-build ok")
    log(f"token ok expires_in={token.expires_in} refresh=yes")
    return token


def sso_to_token(
    sso_cookie: str,
    *,
    proxy: str | None = None,
    log: LogFn | None = None,
    require_referrer: bool = True,
) -> TokenResult:
    """SSO cookie → 带 referrer=grok-build 的 OAuth token。

    仅用 curl_cffi（Chrome TLS）。不再回退 std requests：
    实测 requests 访问 auth.x.ai / accounts.x.ai 常 ReadTimeout，比 curl 更差。
    连接/TLS/超时由上层切到浏览器 PKCE。
    """
    log = log or _noop_log
    sso = normalize_sso_cookie(sso_cookie)
    if not sso:
        raise OAuthCodeError("sso cookie 为空")

    session = _make_session(proxy, prefer="curl")
    try:
        return _run_sso_flow(
            sso, session=session, log=log, require_referrer=require_referrer
        )
    finally:
        try:
            session.close()
        except Exception:
            pass


def mint_from_sso(
    sso_cookie: str,
    *,
    proxy: str | None = None,
    log: LogFn | None = None,
    require_referrer: bool = True,
) -> dict[str, Any]:
    """上层统一返回 dict，兼容 cpa_export。"""
    tr = sso_to_token(
        sso_cookie,
        proxy=proxy,
        log=log,
        require_referrer=require_referrer,
    )
    return {
        "access_token": tr.access_token,
        "refresh_token": tr.refresh_token,
        "id_token": tr.id_token,
        "token_type": tr.token_type,
        "expires_in": tr.expires_in,
        "referrer": tr.referrer,
        "sso": normalize_sso_cookie(sso_cookie),
    }


def _is_http_tls_failure(exc: BaseException | str) -> bool:
    """HTTP 层 TLS/连接/超时失败：适合改走浏览器铸造。"""
    text = str(exc or "").lower()
    needles = (
        "unexpected_eof_while_reading",
        "sslerror",
        "ssleoferror",
        "max retries exceeded",
        "openssl_internal:invalid library",
        "tls connect error",
        "curl: (35)",
        "curl: (28)",
        "failed to perform, curl: (35)",
        "failed to perform, curl: (28)",
        "connection timed out",
        "ssl_error_syscall",
        "connection reset",
        "connection aborted",
        "name resolution",
        "readtimeout",
        "connecttimeout",
        "timed out",
        "timeout",
    )
    return any(n in text for n in needles)


def _page_eval(page: Any, js: str, *args: Any) -> Any:
    """兼容 DrissionPage page.run_js / page.run_js_loaded。"""
    if page is None:
        raise OAuthCodeError("page 为空，无法浏览器铸造")
    last_err: Exception | None = None
    for name in ("run_js", "run_js_loaded", "run_async_js"):
        fn = getattr(page, name, None)
        if not callable(fn):
            continue
        try:
            if args:
                return fn(js, *args)
            return fn(js)
        except TypeError:
            # 某些签名不接受额外参数
            try:
                return fn(js)
            except Exception as exc:  # noqa: BLE001
                last_err = exc
        except Exception as exc:  # noqa: BLE001
            last_err = exc
            continue
    raise OAuthCodeError(f"page 无法执行 JS: {last_err or 'no run_js'}")


def _ensure_sso_on_page(page: Any, sso: str, log: LogFn) -> None:
    sso = normalize_sso_cookie(sso)
    if not sso:
        raise OAuthCodeError("sso cookie 为空")
    # 先落到 accounts 域，再写 cookie，避免 set 失败
    for url in ("https://accounts.x.ai/", "https://auth.x.ai/"):
        try:
            page.get(url)
            time.sleep(0.25)
        except Exception as exc:  # noqa: BLE001
            log(f"open {url} warn: {exc}")

    items = []
    for domain in (".x.ai", "accounts.x.ai", ".accounts.x.ai", "auth.x.ai", ".auth.x.ai"):
        for name in ("sso", "sso-rw"):
            items.append(
                {
                    "name": name,
                    "value": sso,
                    "domain": domain,
                    "path": "/",
                    "secure": True,
                    "sameSite": "None",
                }
            )

    set_ok = False
    # 优先 CDP/DrissionPage set.cookies（可写 httpOnly 域 cookie）
    for target in (page, getattr(page, "browser", None)):
        if target is None:
            continue
        try:
            setter = getattr(target, "set", None)
            cookies_fn = getattr(setter, "cookies", None) if setter is not None else None
            if not callable(cookies_fn):
                continue
            try:
                cookies_fn(items)
                set_ok = True
                log(f"browser sso cookie set bulk via {type(target).__name__}")
                break
            except Exception:
                n = 0
                for it in items:
                    try:
                        cookies_fn(it)
                        n += 1
                    except Exception:
                        pass
                if n:
                    set_ok = True
                    log(f"browser sso cookie set one-by-one={n}")
                    break
        except Exception:
            continue

    # document.cookie 兜底（非 httpOnly）
    set_js = r"""
(sso) => {
  try {
    const maxAge = 60 * 60 * 24 * 30;
    const base = `; path=/; max-age=${maxAge}; SameSite=None; Secure`;
    document.cookie = `sso=${sso}${base}; domain=.x.ai`;
    document.cookie = `sso-rw=${sso}${base}; domain=.x.ai`;
    document.cookie = `sso=${sso}${base}`;
    document.cookie = `sso-rw=${sso}${base}`;
    return document.cookie.includes('sso=');
  } catch (e) {
    return String(e);
  }
}
"""
    try:
        ok = _page_eval(page, set_js, sso)
        log(f"browser sso document.cookie: {ok!r}")
        set_ok = set_ok or (ok is True) or (ok == True) or (str(ok).lower() == "true")
    except Exception as exc:  # noqa: BLE001
        log(f"browser sso document.cookie fail: {exc}")

    if not set_ok:
        raise OAuthCodeError("写入 sso cookie 失败")


def _page_html(page: Any) -> str:
    try:
        html = str(getattr(page, "html", "") or "")
        if html:
            return html
    except Exception:
        pass
    try:
        return str(_page_eval(page, "() => document.documentElement.outerHTML") or "")
    except Exception:
        return ""


def _extract_next_action_id(html: str) -> str:
    """从 consent 页 HTML 抽 Next-Action / Server Action id。"""
    text = html or ""
    patterns = (
        r'"next-action"\s*:\s*"([a-f0-9]{20,})"',
        r'"actionId"\s*:\s*"([a-f0-9]{20,})"',
        r'\$ACTION_ID_([a-f0-9]{20,})',
        r'next-action["\']?\s*[:=]\s*["\']([a-f0-9]{20,})',
    )
    for pat in patterns:
        m = re.search(pat, text, re.I)
        if m:
            return m.group(1)
    return NEXT_ACTION_ID


def _browser_click_allow(page: Any, log: LogFn) -> bool:
    """Consent「允许」必须真实点击（JS click 会导致 Invalid action）。"""
    labels = ("允许", "Allow", "Authorize", "Approve", "授权")
    # 1) DrissionPage 真实点击（与 device consent 同一套经验）
    try:
        candidates = []
        for sel in ("tag:button", "css:button", "css:[role='button']", "css:input[type='submit']"):
            try:
                found = page.eles(sel, timeout=0.3) or []
            except Exception:
                found = []
            for el in found:
                try:
                    t = (getattr(el, "text", None) or el.raw_text or "").strip()
                except Exception:
                    t = ""
                if not t:
                    continue
                compact = re.sub(r"\s+", "", t)
                if compact in labels or t in labels:
                    candidates.append((0, el, t))
                elif any(x in compact for x in labels) and "全部" not in compact and "All" not in compact:
                    candidates.append((1, el, t))
        candidates.sort(key=lambda x: x[0])
        for _, el, t in candidates:
            try:
                el.click()  # real click
                log(f"browser REAL click allow: {t!r}")
                return True
            except Exception as exc:  # noqa: BLE001
                log(f"real click {t!r} failed: {exc}")
                try:
                    el.click(by_js=True)
                    log(f"browser JS click allow: {t!r}")
                    return True
                except Exception:
                    continue
    except Exception as exc:  # noqa: BLE001
        log(f"ele allow click failed: {exc}")

    # 2) 表单 action=allow 后 submit（device consent 兜底同款）
    try:
        ret = _page_eval(
            page,
            r"""
() => {
  const labels = new Set(['允许','Allow','Authorize','Approve','授权']);
  const f = document.querySelector('form');
  if (f) {
    let a = f.querySelector('input[name=action]');
    if (!a) {
      a = document.createElement('input');
      a.type = 'hidden';
      a.name = 'action';
      f.appendChild(a);
    }
    a.value = 'allow';
    const btn = [...f.querySelectorAll('button,[role=button],input[type=submit]')]
      .find(b => labels.has(((b.innerText||b.value||'').trim())));
    if (btn) { btn.click(); return {ok:true, via:'form-btn'}; }
    f.submit();
    return {ok:true, via:'form-submit'};
  }
  // 无 form：真实派发 pointer/mouse 事件
  const nodes = [...document.querySelectorAll('button,[role=button],input[type=submit],a')];
  const target = nodes.find(n => {
    const t = ((n.innerText||n.textContent||n.value||'').replace(/\s+/g,'')).trim();
    return labels.has(t) || (t.includes('允许') && !t.includes('全部'));
  });
  if (!target) {
    return {ok:false, texts: nodes.slice(0,10).map(n => (n.innerText||n.value||'').trim()).filter(Boolean)};
  }
  target.scrollIntoView({block:'center'});
  target.focus();
  for (const type of ['pointerdown','mousedown','pointerup','mouseup','click']) {
    target.dispatchEvent(new MouseEvent(type, {bubbles:true, cancelable:true, view:window}));
  }
  return {ok:true, via:'mouse-events', text:(target.innerText||target.value||'').trim()};
}
""",
        )
        if isinstance(ret, dict) and ret.get("ok"):
            log(f"browser allow fallback: {ret}")
            return True
        log(f"browser allow button not found: {ret!r}")
    except Exception as exc:  # noqa: BLE001
        log(f"click allow js failed: {exc}")
    return False


def _browser_consent_server_action(
    page: Any,
    consent_url: str,
    flow: AuthCodeFlow,
    log: LogFn,
) -> str:
    """在浏览器内 POST Next.js Server Action 批准 consent，直接拿 code。"""
    html = _page_html(page)
    action_id = _extract_next_action_id(html)
    payload = [
        {
            "action": "allow",
            "clientId": CLIENT_ID,
            "redirectUri": REDIRECT_URI,
            "scope": SCOPE,
            "state": flow.state,
            "codeChallenge": flow.code_challenge,
            "codeChallengeMethod": "S256",
            "nonce": flow.nonce,
            "principalType": "User",
            "principalId": "",
            "referrer": GROK_REFERRER,
        }
    ]
    body = json.dumps(payload, separators=(",", ":"))
    # 用 window key 轮询，兼容 DrissionPage 对 Promise 支持不一
    wrap = f"""
(() => {{
  const key = '__cpa_consent_' + Date.now();
  window[key] = null;
  fetch({consent_url!r}, {{
    method: 'POST',
    headers: {{
      'Accept': 'text/x-component',
      'Content-Type': 'text/plain;charset=UTF-8',
      'Next-Action': {action_id!r},
      'Origin': 'https://accounts.x.ai',
      'Referer': {consent_url!r},
    }},
    body: {body!r},
    credentials: 'include',
    redirect: 'follow',
  }}).then(async (r) => {{
    const text = await r.text();
    window[key] = {{status: r.status, url: r.url, text: text}};
  }}).catch((e) => {{
    window[key] = {{status: 0, url: '', text: String(e)}};
  }});
  return key;
}})()
"""
    try:
        key = _page_eval(page, wrap)
    except Exception as exc:  # noqa: BLE001
        raise OAuthCodeError(f"browser consent fetch 启动失败: {exc}") from exc
    log(f"browser consent server-action post action_id={action_id[:12]}…")
    deadline = time.time() + 20
    ret = None
    while time.time() < deadline:
        try:
            ret = _page_eval(page, f"() => window[{key!r}]")
        except Exception:
            try:
                ret = _page_eval(page, f"window[{key!r}]")
            except Exception:
                ret = None
        if ret:
            break
        time.sleep(0.2)
    if not isinstance(ret, dict):
        raise OAuthCodeError(f"browser consent 无响应: {ret!r}")
    status = int(ret.get("status") or 0)
    text = str(ret.get("text") or "")
    final = str(ret.get("url") or "")
    if status and (status < 200 or status >= 300):
        raise OAuthCodeError(f"browser consent HTTP {status}: {_short(text, 240)}")
    if "code=" in final:
        qs = parse_qs(urlparse(final).query)
        code = (qs.get("code") or [""])[0].strip()
        if code:
            log("browser got code from consent redirect url")
            return code
    # 响应体 / RSC
    try:
        code = parse_consent_code(text)
        if code:
            log("browser got code from consent server-action body")
            return code
    except OAuthCodeError:
        pass
    # 有时 code 在 text 的 redirect 串里
    m = re.search(r"[?&]code=([A-Za-z0-9._~\-]+)", text)
    if m:
        log("browser got code from consent body regex")
        return m.group(1)
    raise OAuthCodeError(f"browser consent 未返回 code: {_short(text, 240)}")


def _browser_fetch_token(page: Any, code: str, flow: AuthCodeFlow, log: LogFn) -> TokenResult:
    """在浏览器上下文用 fetch 换 token，绕开 Python TLS 对 auth.x.ai 的 EOF。"""
    form = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": REDIRECT_URI,
        "client_id": CLIENT_ID,
        "code_verifier": flow.code_verifier,
    }
    body = urlencode(form)
    js = r"""
(tokenUrl, body, ua, ver) => {
  return fetch(tokenUrl, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/x-www-form-urlencoded',
      'Accept': '*/*',
      'User-Agent': ua,
      'X-Grok-Client-Version': ver,
    },
    body: body,
    credentials: 'include',
  }).then(async (r) => {
    const text = await r.text();
    return {status: r.status, text: text};
  }).catch((e) => ({status: 0, text: String(e)}));
}
"""
    # 先到 auth 域，减少跨站限制
    try:
        page.get(ISSUER + "/")
        time.sleep(0.3)
    except Exception:
        pass
    ret = None
    # DrissionPage 对 Promise 支持不一，做短轮询包装
    wrap = r"""
(tokenUrl, body, ua, ver) => {
  const key = '__cpa_token_result_' + Date.now();
  window[key] = null;
  fetch(tokenUrl, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/x-www-form-urlencoded',
      'Accept': '*/*',
      'User-Agent': ua,
      'X-Grok-Client-Version': ver,
    },
    body: body,
    credentials: 'include',
  }).then(async (r) => {
    const text = await r.text();
    window[key] = {status: r.status, text: text};
  }).catch((e) => {
    window[key] = {status: 0, text: String(e)};
  });
  return key;
}
"""
    try:
        key = _page_eval(page, wrap, TOKEN_URL, body, GROK_TOKEN_UA, GROK_VERSION)
    except Exception:
        # 无参回退：把参数内联
        key = _page_eval(
            page,
            f"""
(() => {{
  const key = '__cpa_token_result_' + Date.now();
  window[key] = null;
  fetch({TOKEN_URL!r}, {{
    method: 'POST',
    headers: {{
      'Content-Type': 'application/x-www-form-urlencoded',
      'Accept': '*/*',
      'User-Agent': {GROK_TOKEN_UA!r},
      'X-Grok-Client-Version': {GROK_VERSION!r},
    }},
    body: {body!r},
    credentials: 'include',
  }}).then(async (r) => {{
    const text = await r.text();
    window[key] = {{status: r.status, text: text}};
  }}).catch((e) => {{
    window[key] = {{status: 0, text: String(e)}};
  }});
  return key;
}})()
""",
        )
    deadline = time.time() + 30
    while time.time() < deadline:
        try:
            ret = _page_eval(page, f"() => window[{key!r}]")
        except Exception:
            try:
                ret = _page_eval(page, f"window[{key!r}]")
            except Exception as exc:
                raise OAuthCodeError(f"读取 browser token 结果失败: {exc}") from exc
        if ret:
            break
        time.sleep(0.2)
    if not isinstance(ret, dict):
        raise OAuthCodeError(f"browser token 无响应: {ret!r}")
    status = int(ret.get("status") or 0)
    text = str(ret.get("text") or "")
    if status < 200 or status >= 300:
        raise OAuthCodeError(f"browser token HTTP {status}: {_short(text, 300)}")
    try:
        data = json.loads(text)
    except Exception as e:
        raise OAuthCodeError(f"browser token 非 JSON: {_short(text)}") from e
    if not isinstance(data, dict) or not data.get("access_token"):
        raise OAuthCodeError(f"browser token 缺少 access_token: {data!r}")
    access = str(data["access_token"]).strip()
    refresh = str(data.get("refresh_token") or "").strip()
    if not refresh:
        raise OAuthCodeError("browser token 缺少 refresh_token")
    referrer = ""
    try:
        referrer = str(jwt_payload(access).get("referrer") or "")
    except Exception:
        pass
    return TokenResult(
        access_token=access,
        refresh_token=refresh,
        id_token=(str(data["id_token"]).strip() if data.get("id_token") else None),
        token_type=str(data.get("token_type") or "Bearer"),
        expires_in=int(data.get("expires_in") or 21600),
        raw=data,
        referrer=referrer,
    )


def mint_from_sso_browser(
    sso_cookie: str,
    page: Any,
    *,
    log: LogFn | None = None,
    require_referrer: bool = True,
    timeout_sec: float = 90.0,
) -> dict[str, Any]:
    """用注册浏览器完成 SSO→PKCE（绕开 Python TLS 访问 auth.x.ai 失败）。

    流程：
      1. 写入 sso cookie
      2. 打开 authorize(referrer=grok-build)
      3. 在 consent 页点击允许 / 或解析 callback code
      4. 浏览器 fetch oauth2/token
    """
    log = log or _noop_log
    sso = normalize_sso_cookie(sso_cookie)
    if not sso:
        raise OAuthCodeError("sso cookie 为空")
    if page is None:
        raise OAuthCodeError("page 为空")

    flow = new_auth_code_flow()
    params = {
        "response_type": "code",
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "scope": SCOPE,
        "code_challenge": flow.code_challenge,
        "code_challenge_method": "S256",
        "state": flow.state,
        "nonce": flow.nonce,
        "referrer": GROK_REFERRER,
    }
    auth_url = f"{AUTHORIZE_URL}?{urlencode(params)}"
    log(f"browser PKCE authorize referrer={GROK_REFERRER}")
    _ensure_sso_on_page(page, sso, log)

    try:
        page.get(auth_url)
    except Exception as exc:  # noqa: BLE001
        raise OAuthCodeError(f"browser 打开 authorize 失败: {exc}") from exc

    code = ""
    deadline = time.time() + max(30.0, float(timeout_sec))
    last_url = ""
    consent_attempts = 0
    server_action_tried = False
    wait_after_click_until = 0.0

    def _code_from_url(u: str) -> str:
        if not u or "code=" not in u:
            return ""
        if not any(x in u for x in ("127.0.0.1", "localhost", "callback", "redirect")):
            # 仍允许任意带 code 的 redirect
            if "code=" not in u:
                return ""
        qs = parse_qs(urlparse(u).query)
        return (qs.get("code") or [""])[0].strip()

    while time.time() < deadline:
        try:
            url = str(getattr(page, "url", "") or "")
        except Exception:
            url = ""
        if url and url != last_url:
            log(f"browser url: {_short(url, 140)}")
            last_url = url

        # callback 已跳到 redirect_uri?code=
        code = _code_from_url(url)
        if code:
            log("browser got code from redirect")
            break

        # consent 页：真实点击 Allow；点后短暂等待跳转；仍不行再 Server Action POST
        if "/oauth2/consent" in url or "/consent" in url:
            # 刚点过，先等 redirect
            if wait_after_click_until and time.time() < wait_after_click_until:
                time.sleep(0.3)
                continue

            consent_attempts += 1
            if consent_attempts <= 4:
                clicked = _browser_click_allow(page, log)
                if clicked:
                    wait_after_click_until = time.time() + 4.0
                    time.sleep(0.6)
                    # 点后立刻看 url / html
                    try:
                        url2 = str(getattr(page, "url", "") or "")
                    except Exception:
                        url2 = url
                    code = _code_from_url(url2)
                    if code:
                        log("browser got code after allow click")
                        break
                    html = _page_html(page)
                    if html:
                        try:
                            code = parse_consent_code(html)
                            if code:
                                log("browser got code from consent html after click")
                                break
                        except OAuthCodeError:
                            pass
                    continue

            # 真实点击无效：浏览器内 Server Action POST（与 HTTP approve 同 payload）
            if not server_action_tried:
                server_action_tried = True
                try:
                    code = _browser_consent_server_action(page, url, flow, log)
                    if code:
                        break
                except Exception as exc:  # noqa: BLE001
                    log(f"browser consent server-action failed: {exc}")
            time.sleep(0.5)
            continue

        if "sign-in" in url or "sign-up" in url:
            # cookie 可能没带上，重写一次
            log("browser landed sign-in, re-inject sso")
            _ensure_sso_on_page(page, sso, log)
            try:
                page.get(auth_url)
            except Exception:
                pass
            time.sleep(0.8)
            continue

        # 中间跳转页：稍等
        time.sleep(0.4)

    if not code:
        # 最后再扫一次页面 HTML
        try:
            html = _page_html(page)
            if html:
                try:
                    code = parse_consent_code(html)
                except OAuthCodeError:
                    code = ""
        except Exception:
            code = ""
    if not code:
        raise OAuthCodeError(
            f"browser PKCE 超时未拿到 code（last_url={_short(last_url, 160)}）"
        )

    token = _browser_fetch_token(page, code, flow, log)
    if token.referrer != GROK_REFERRER:
        msg = f"access_token 未包含预期 referrer（got={token.referrer!r}）"
        if require_referrer:
            raise OAuthCodeError(msg)
        log(f"WARN {msg}")
    else:
        log("browser access_token referrer=grok-build ok")
    log(f"browser token ok expires_in={token.expires_in}")
    return {
        "access_token": token.access_token,
        "refresh_token": token.refresh_token,
        "id_token": token.id_token,
        "token_type": token.token_type,
        "expires_in": token.expires_in,
        "referrer": token.referrer,
        "sso": sso,
    }