"""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 page
        -> POST auth.x.ai/oauth2/authorize (form action=allow)
        -> oauth2/token (authorization_code + PKCE)

2026-07-16：accounts.x.ai consent 页 Next.js Server Action 频繁轮换，
硬编码 Next-Action 会 404 "Server action not found"。
官方 HTML form 的 action 指向 auth.x.ai/oauth2/authorize（非 consent URL），
form-urlencoded + action=allow 可稳定拿到 code。
"""

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"
# 官方 CLI 0.2.101：loopback redirect，运行时随机端口（RFC 8252 端口无关）
# 兼容保留旧固定端口常量，仅作 fallback
REDIRECT_URI_LEGACY = "http://127.0.0.1:56121/callback"
REDIRECT_URI = REDIRECT_URI_LEGACY
# cli-chat-proxy 要求 access_token.scope 含 grok-cli:access，否则：
#   WKE=unauthorized:grok-cli-token-auth-required
# 2026-07 一度误删该 scope（对齐 discovery），导致铸造成功但调用 403。
SCOPE = (
    "openid profile email offline_access "
    "grok-cli:access api:access conversations:read conversations:write"
)
GROK_CLI_SCOPE = "grok-cli:access"
GROK_REFERRER = "grok-build"
# 对齐官方 stable CLI（2026-07-14 二进制：0.2.101）
GROK_VERSION = "0.2.101"
GROK_TOKEN_UA = f"xai-grok-build/{GROK_VERSION}"
GROK_CLIENT_SURFACE = "grok-build"
# 旧 Next.js Server Action id（已失效，仅作 fallback 尝试）
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
    redirect_uri: str = REDIRECT_URI_LEGACY


@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_loopback_redirect_uri() -> str:
    """官方 CLI：http://127.0.0.1:<ephemeral>/callback。"""
    # 高位端口，避免与常见本机服务冲突
    port = 49152 + secrets.randbelow(16383)
    return f"http://127.0.0.1:{port}/callback"


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,
        redirect_uri=_new_loopback_redirect_uri(),
    )


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 is_cloudflare_block(
    status: int | None = None,
    body: str = "",
    exc: BaseException | str | None = None,
) -> bool:
    """识别 auth.x.ai 被 Cloudflare 拦（常见直连 403 挑战页）。"""
    text = f"{body or ''} {exc or ''}".lower()
    if status == 403 and (
        "<!doctype html" in text
        or "cloudflare" in text
        or "cf-ray" in text
        or "attention required" in text
        or "just a moment" in text
        or "enable javascript" in text
        or "oldie" in text
    ):
        return True
    if "authorize http 403" in text and (
        "<!doctype" in text or "oldie" in text or "cloudflare" in text
    ):
        return True
    return "cloudflare" in text and ("403" in text or "blocked" in text)


def _cf_block_hint(proxy_label: str = "") -> str:
    via = proxy_label or "(direct)"
    return (
        f"Cloudflare 拦截 auth.x.ai（出口={via}）。"
        "直连大陆/机房 IP 几乎必 403；请配置 mint_proxy / proxy，"
        "或开启 proxy_pool_enabled 并保证代理能访问 auth.x.ai。"
    )


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]:
    # 对齐官方 0.2.101：小写 x-grok-* + surface=grok-build
    return {
        "User-Agent": GROK_TOKEN_UA,
        "Accept": "*/*",
        "Content-Type": "application/x-www-form-urlencoded",
        "x-grok-client-version": GROK_VERSION,
        "x-grok-client-surface": GROK_CLIENT_SURFACE,
        # 兼容旧中间件/代理仍读 Pascal 头
        "X-Grok-Client-Version": GROK_VERSION,
    }


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


def open_authorize_page(session: Any, flow: AuthCodeFlow) -> str:
    redirect_uri = flow.redirect_uri or _new_loopback_redirect_uri()
    flow.redirect_uri = redirect_uri
    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=30,
    )
    body = resp.text or ""
    final = _final_url(resp)
    if resp.status_code < 200 or resp.status_code >= 300:
        if is_cloudflare_block(resp.status_code, body):
            raise OAuthCodeError(
                f"authorize HTTP {resp.status_code}: Cloudflare 拦截 — {_short(body, 80)}"
            )
        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 _code_from_location(url: str) -> str:
    if not url or "code=" not in url:
        return ""
    qs = parse_qs(urlparse(url).query)
    return str((qs.get("code") or [""])[0] or "").strip()


def _approve_via_form_post(
    session: Any, consent_url: str, flow: AuthCodeFlow
) -> str:
    """POST form to auth.x.ai/oauth2/authorize (matches consent page HTML)."""
    redirect_uri = flow.redirect_uri or REDIRECT_URI_LEGACY
    fields = {
        "client_id": CLIENT_ID,
        "redirect_uri": redirect_uri,
        "scope": SCOPE,
        "state": flow.state,
        "code_challenge": flow.code_challenge,
        "code_challenge_method": "S256",
        "nonce": flow.nonce,
        "principal_type": "User",
        "principal_id": "",
        "referrer": GROK_REFERRER,
        "action": "allow",
    }
    headers = {
        "User-Agent": BROWSER_UA,
        "Accept": (
            "text/html,application/xhtml+xml,application/xml;q=0.9,"
            "*/*;q=0.8"
        ),
        "Content-Type": "application/x-www-form-urlencoded",
        "Origin": "https://accounts.x.ai",
        "Referer": consent_url,
        "Sec-Fetch-Site": "same-site",
        "Sec-Fetch-Mode": "navigate",
        "Sec-Fetch-Dest": "document",
        "Upgrade-Insecure-Requests": "1",
        "Accept-Language": "en-US,en;q=0.9",
    }
    # 不自动 follow 到 127.0.0.1 loopback（本机无 listener）
    resp = session.post(
        AUTHORIZE_URL,
        data=urlencode(fields),
        headers=headers,
        allow_redirects=False,
        timeout=30,
    )
    loc = (
        resp.headers.get("Location")
        or resp.headers.get("location")
        or ""
    )
    code = _code_from_location(loc)
    if code:
        return code

    # 少数库会吞 Location 到 resp.url / 已 follow
    final = _final_url(resp)
    code = _code_from_location(final)
    if code:
        return code

    text = resp.text or ""
    if 200 <= resp.status_code < 300:
        try:
            return parse_consent_code(text)
        except OAuthCodeError:
            pass
    raise OAuthCodeError(
        f"consent form POST HTTP {resp.status_code}: "
        f"loc={_short(loc, 120)} body={_short(text, 200)}"
    )


def _approve_via_next_action(
    session: Any, consent_url: str, flow: AuthCodeFlow
) -> str:
    """旧路径：Next.js Server Action POST consent URL（action id 常失效）。"""
    redirect_uri = flow.redirect_uri or REDIRECT_URI_LEGACY
    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=False,
        timeout=30,
    )
    text = resp.text or ""
    loc = (
        resp.headers.get("Location")
        or resp.headers.get("location")
        or ""
    )
    code = _code_from_location(loc) or _code_from_location(_final_url(resp))
    if code:
        return code
    if resp.status_code < 200 or resp.status_code >= 300:
        raise OAuthCodeError(
            f"consent Next-Action HTTP {resp.status_code}: {_short(text, 300)}"
        )
    return parse_consent_code(text)


def approve_authorization(session: Any, consent_url: str, flow: AuthCodeFlow) -> str:
    """Approve consent and return authorization code.

    优先：HTML form POST → https://auth.x.ai/oauth2/authorize
    回退：旧 Next-Action POST consent URL（易 404）。
    """
    try:
        return _approve_via_form_post(session, consent_url, flow)
    except OAuthCodeError as form_exc:
        try:
            return _approve_via_next_action(session, consent_url, flow)
        except OAuthCodeError as next_exc:
            raise OAuthCodeError(
                f"consent allow failed: form={form_exc}; next_action={next_exc}"
            ) from next_exc


def exchange_auth_code(session: Any, code: str, flow: AuthCodeFlow) -> TokenResult:
    redirect_uri = flow.redirect_uri or REDIRECT_URI_LEGACY
    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=30,
    )
    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 ver={GROK_VERSION} ua={GROK_TOKEN_UA} "
        f"referrer={GROK_REFERRER} redirect={flow.redirect_uri} 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")

    # cli-chat-proxy 的 Grok CLI gate：缺 grok-cli:access 会直接 403
    scope_text = ""
    try:
        scope_text = str(jwt_payload(token.access_token).get("scope") or "")
    except Exception:
        scope_text = ""
    scopes = set(scope_text.split())
    if GROK_CLI_SCOPE not in scopes:
        msg = (
            f"access_token 缺少 {GROK_CLI_SCOPE} "
            f"(scope={scope_text or '(empty)'})；"
            "cli-chat-proxy 会返回 grok-cli-token-auth-required"
        )
        if require_referrer:
            raise OAuthCodeError(msg)
        log(f"WARN {msg}")
    else:
        log(f"access_token scope 含 {GROK_CLI_SCOPE} 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。"""
    from .proxyutil import proxy_log_label, resolve_proxy

    log = log or _noop_log
    sso = normalize_sso_cookie(sso_cookie)
    if not sso:
        raise OAuthCodeError("sso cookie 为空")

    resolved = resolve_proxy(proxy)
    log(f"mint 出口={proxy_log_label(resolved) or '(direct)'}")

    # 先 curl_cffi；若遇到 OpenSSL invalid library / curl(35)，自动回退 std requests
    session = _make_session(proxy, prefer="curl")
    try:
        return _run_sso_flow(
            sso, session=session, log=log, require_referrer=require_referrer
        )
    except Exception as exc:  # noqa: BLE001
        backend = str(getattr(session, "_cpa_http_backend", "") or "")
        can_fallback = _is_curl_tls_broken(exc) or (
            backend.startswith("curl_cffi") and "curl: (35)" in str(exc).lower()
        )
        if not can_fallback:
            if is_cloudflare_block(exc=exc):
                raise OAuthCodeError(_cf_block_hint(proxy_log_label(resolved))) from exc
            raise
        log(f"curl TLS 异常，回退标准 requests: {_short(str(exc), 160)}")
        try:
            session.close()
        except Exception:
            pass
        session = _make_session(proxy, prefer="requests")
        try:
            return _run_sso_flow(
                sso, session=session, log=log, require_referrer=require_referrer
            )
        except Exception as exc2:  # noqa: BLE001
            if is_cloudflare_block(exc=exc2):
                raise OAuthCodeError(_cf_block_hint(proxy_log_label(resolved))) from exc2
            raise
        finally:
            try:
                session.close()
            except Exception:
                pass
            session = None  # type: ignore[assignment]
    finally:
        if session is not None:
            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),
    }