"""CPA 上传：校验 session.account.planType=='plus'、构造 codex auth JSON、POST 到 CPA。

参考 chrome_extension/codex-oauth-automation-extension/background/cpa-api.js。
"""
from __future__ import annotations

import base64
import json
import re
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from typing import Callable
from urllib.parse import urlparse, quote


def _normalize_str(value) -> str:
    return str(value or "").strip()


def _is_email(value: str) -> bool:
    return bool(value and re.match(r"^[^\s@]+@[^\s@]+\.[^\s@]+$", value))


def _first_non_empty(*values) -> str:
    for v in values:
        s = _normalize_str(v)
        if s:
            return s
    return ""


def _b64url_decode(segment: str) -> str:
    s = _normalize_str(segment).replace("-", "+").replace("_", "/")
    if not s:
        return ""
    pad = (-len(s)) % 4
    s += "=" * pad
    try:
        return base64.b64decode(s).decode("utf-8", errors="replace")
    except Exception:
        return ""


def _b64url_encode_json(value) -> str:
    raw = json.dumps(value, separators=(",", ":")).encode("utf-8")
    enc = base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
    return enc


def parse_jwt_payload(token: str) -> dict:
    token = _normalize_str(token)
    if not token:
        return {}
    parts = token.split(".")
    if len(parts) < 2:
        return {}
    decoded = _b64url_decode(parts[1])
    if not decoded:
        return {}
    try:
        return json.loads(decoded)
    except Exception:
        return {}


def _normalize_iso_timestamp(value) -> str:
    if isinstance(value, datetime):
        return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
    if isinstance(value, (int, float)):
        ms = value if value > 1e11 else value * 1000
        try:
            return datetime.fromtimestamp(ms / 1000, tz=timezone.utc).isoformat().replace("+00:00", "Z")
        except Exception:
            return ""
    if isinstance(value, str) and value.strip():
        try:
            dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
            return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
        except Exception:
            return ""
    return ""


def _ts_from_unix_seconds(value) -> str:
    try:
        n = float(value)
    except Exception:
        return ""
    try:
        return datetime.fromtimestamp(n, tz=timezone.utc).isoformat().replace("+00:00", "Z")
    except Exception:
        return ""


def _epoch_seconds_from(value) -> int:
    if value in (None, ""):
        return 0
    try:
        n = float(value)
        return int(n / 1000) if n > 1e11 else int(n)
    except Exception:
        pass
    iso = _normalize_iso_timestamp(value)
    if not iso:
        return 0
    try:
        return int(datetime.fromisoformat(iso.replace("Z", "+00:00")).timestamp())
    except Exception:
        return 0


def _build_synthetic_id_token(email, account_id, plan_type, user_id, expires_at) -> str:
    if not _normalize_str(account_id):
        return ""
    now = int(time.time())
    expires = _epoch_seconds_from(expires_at) or (now + 90 * 24 * 60 * 60)
    auth_info = {"chatgpt_account_id": _normalize_str(account_id)}
    if plan_type:
        auth_info["chatgpt_plan_type"] = _normalize_str(plan_type)
    if user_id:
        auth_info["chatgpt_user_id"] = _normalize_str(user_id)
        auth_info["user_id"] = _normalize_str(user_id)
    payload = {
        "iat": now,
        "exp": expires,
        "https://api.openai.com/auth": auth_info,
    }
    if email:
        payload["email"] = _normalize_str(email)
    header = _b64url_encode_json({"alg": "none", "typ": "JWT", "cpa_synthetic": True})
    body = _b64url_encode_json(payload)
    return f"{header}.{body}.synthetic"


def _sanitize_segment(value: str, fallback: str = "chatgpt-session") -> str:
    s = _normalize_str(value)
    s = re.sub(r"[\\/:*?\"<>|]+", "-", s)
    s = re.sub(r"\s+", "-", s)
    s = re.sub(r"-+", "-", s).strip("-")
    return s or fallback


def _normalize_plan_for_filename(plan_type: str) -> str:
    parts = re.split(r"[^a-zA-Z0-9]+", _normalize_str(plan_type))
    return "-".join(p.lower() for p in parts if p)


def build_cpa_filename(email: str, plan_type: str, account_id: str) -> str:
    e = _sanitize_segment(email or "")
    p = _normalize_plan_for_filename(plan_type or "")
    a = _sanitize_segment(account_id or "")
    if e and p:
        return f"codex-{e}-{p}.json"
    if e:
        return f"codex-{e}.json"
    if a and p:
        return f"codex-{a}-{p}.json"
    if a:
        return f"codex-{a}.json"
    return f"codex-{int(time.time()*1000)}.json"


def get_session_plan_type(session: dict) -> str:
    """读 session.account.planType（与扩展一致）。"""
    if not isinstance(session, dict):
        return ""
    account = session.get("account") or {}
    if not isinstance(account, dict):
        return ""
    return _first_non_empty(account.get("planType"), account.get("plan_type"))


def is_plus_session(session: dict) -> bool:
    plan = get_session_plan_type(session).lower()
    return plan == "plus"


def build_cpa_auth_payload(session: dict, *, email_hint: str = "") -> dict:
    """从 ChatGPT /api/auth/session JSON 构造 CPA codex auth JSON。
    返回 { authJson, accountId, email, fileName, hasRefreshToken }。
    """
    if not isinstance(session, dict):
        raise RuntimeError("session 不是 JSON 对象")

    access_token = _normalize_str(session.get("accessToken"))
    if not access_token:
        raise RuntimeError("session 中没有 accessToken")

    input_id_token = _first_non_empty(session.get("idToken"), session.get("id_token"))
    refresh_token = _first_non_empty(session.get("refreshToken"), session.get("refresh_token"))
    session_token = _first_non_empty(session.get("sessionToken"), session.get("session_token"))

    access_payload = parse_jwt_payload(access_token)
    id_payload = parse_jwt_payload(input_id_token)
    access_auth = access_payload.get("https://api.openai.com/auth") if isinstance(access_payload, dict) else {}
    id_auth = id_payload.get("https://api.openai.com/auth") if isinstance(id_payload, dict) else {}
    profile = access_payload.get("https://api.openai.com/profile") if isinstance(access_payload, dict) else {}
    if not isinstance(access_auth, dict):
        access_auth = {}
    if not isinstance(id_auth, dict):
        id_auth = {}
    if not isinstance(profile, dict):
        profile = {}

    expires_at = _first_non_empty(
        _ts_from_unix_seconds(access_payload.get("exp")) if isinstance(access_payload, dict) else "",
        _normalize_iso_timestamp(session.get("expires")),
        _normalize_iso_timestamp(session.get("expiresAt")),
        _normalize_iso_timestamp(session.get("expired")),
        _normalize_iso_timestamp(session.get("expires_at")),
    )

    user = session.get("user") if isinstance(session.get("user"), dict) else {}
    account = session.get("account") if isinstance(session.get("account"), dict) else {}

    def _email(v):
        s = _normalize_str(v).lower()
        return s if _is_email(s) else ""

    email = _first_non_empty(
        _email(user.get("email")),
        _email(session.get("email")),
        _email(email_hint),
        _email(profile.get("email")) if profile else "",
        _email(id_payload.get("email")) if isinstance(id_payload, dict) else "",
        _email(access_payload.get("email")) if isinstance(access_payload, dict) else "",
    )

    account_id = _first_non_empty(
        account.get("id") if account else "",
        session.get("account_id"),
        access_auth.get("chatgpt_account_id"),
        id_auth.get("chatgpt_account_id"),
    )

    user_id = _first_non_empty(
        user.get("id") if user else "",
        session.get("user_id"),
        access_auth.get("chatgpt_user_id"),
        access_auth.get("user_id"),
        id_auth.get("chatgpt_user_id"),
        id_auth.get("user_id"),
    )

    plan_type = _first_non_empty(
        account.get("planType") if account else "",
        account.get("plan_type") if account else "",
        session.get("planType"),
        session.get("plan_type"),
        access_auth.get("chatgpt_plan_type"),
        id_auth.get("chatgpt_plan_type"),
    )

    exported_at = _normalize_iso_timestamp(datetime.now(timezone.utc)) or datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    synthetic_id_token = "" if input_id_token else _build_synthetic_id_token(email, account_id, plan_type, user_id, expires_at)
    id_token = input_id_token or synthetic_id_token

    auth_json_full = {
        "type": "codex",
        "account_id": account_id,
        "chatgpt_account_id": account_id,
        "email": email,
        "name": _first_non_empty(email, "ChatGPT Account"),
        "plan_type": plan_type,
        "chatgpt_plan_type": plan_type,
        "id_token": id_token,
        "id_token_synthetic": True if synthetic_id_token else None,
        "access_token": access_token,
        "refresh_token": refresh_token or "",
        "session_token": session_token,
        "last_refresh": exported_at,
        "expired": expires_at,
        "disabled": True if session.get("disabled") is True else None,
    }
    auth_json = {k: v for k, v in auth_json_full.items() if v not in (None, "")}

    return {
        "authJson": auth_json,
        "accountId": account_id,
        "email": email,
        "expiresAt": expires_at,
        "fileName": build_cpa_filename(email, plan_type, account_id),
        "hasRefreshToken": bool(refresh_token),
        "planType": plan_type,
    }


def _http_post_json(url: str, *, headers: dict, body: dict, timeout: int = 60) -> tuple[int, str, dict]:
    """POST JSON。优先 curl_cffi 走 chrome 指纹（绕 Cloudflare 等基于 UA 的拦截），失败回退 urllib。"""
    payload_bytes = json.dumps(body).encode("utf-8")

    try:
        from curl_cffi import requests as curl_requests
    except Exception:
        curl_requests = None

    if curl_requests is not None:
        try:
            r = curl_requests.post(
                url,
                data=payload_bytes,
                headers=headers,
                impersonate="chrome136",
                timeout=timeout,
            )
            text = r.text
            status = r.status_code
            try:
                parsed = r.json() if text else {}
            except Exception:
                parsed = {}
            return status, text, parsed
        except Exception:
            # 回退 urllib
            pass

    req = urllib.request.Request(url, data=payload_bytes, method="POST")
    for k, v in headers.items():
        req.add_header(k, v)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            text = resp.read().decode("utf-8", errors="replace")
            status = resp.status
    except urllib.error.HTTPError as exc:
        try:
            text = exc.read().decode("utf-8", errors="replace")
        except Exception:
            text = ""
        status = exc.code
    parsed = {}
    try:
        parsed = json.loads(text or "{}")
    except Exception:
        parsed = {}
    return status, text, parsed


def upload_session_to_cpa(
    session: dict,
    *,
    cpa_url: str,
    management_key: str,
    email_hint: str = "",
    timeout: int = 60,
    log: Callable[[str], None] = print,
) -> dict:
    """把 ChatGPT session 通过 CPA 管理接口上传。

    Returns: { fileName, email, planType, hasRefreshToken, status, response }
    """
    cpa_url = _normalize_str(cpa_url)
    management_key = _normalize_str(management_key)
    if not cpa_url:
        raise RuntimeError("CPA 地址未配置")
    if not management_key:
        raise RuntimeError("CPA 管理密钥未配置")

    parsed_url = urlparse(cpa_url)
    if not parsed_url.scheme or not parsed_url.netloc:
        raise RuntimeError(f"CPA 地址格式无效: {cpa_url}")
    origin = f"{parsed_url.scheme}://{parsed_url.netloc}"

    payload = build_cpa_auth_payload(session, email_hint=email_hint)
    log(f"[cpa] 构造 auth JSON 完成 file={payload['fileName']} email={payload['email']!r} plan={payload['planType']!r} hasRefreshToken={payload['hasRefreshToken']}")
    if not payload["hasRefreshToken"]:
        log("[cpa] 警告：缺少 refresh_token，access_token 过期后无法续期")

    name_q = quote(payload["fileName"], safe="")
    url = f"{origin}/v0/management/auth-files?name={name_q}"
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {management_key}",
        "X-Management-Key": management_key,
    }
    log(f"[cpa] POST {url}")
    status, text, resp = _http_post_json(url, headers=headers, body=payload["authJson"], timeout=timeout)
    log(f"[cpa] HTTP {status} 返回长度 {len(text)} 预览={text[:300]}")

    if status >= 400:
        msg = ""
        if isinstance(resp, dict):
            for k in ("error", "message", "detail", "reason"):
                if resp.get(k):
                    msg = str(resp[k])
                    break
        raise RuntimeError(f"CPA 上传失败 HTTP {status}: {msg or text[:300]}")

    return {
        "fileName": payload["fileName"],
        "email": payload["email"],
        "planType": payload["planType"],
        "hasRefreshToken": payload["hasRefreshToken"],
        "status": status,
        "response": resp,
    }