"""Register-machine hook: mint CPA xai auth after successful registration.

OIDC package lives at ./cpa_xai (bundled with this project).
Optional override: config `api_reverse_tools` / env `API_REVERSE_TOOLS`
points at a directory that *contains* the `cpa_xai` package.

After mint, optionally:
  - copy into local CPA auth-dir (`cpa_copy_to_hotload`)
  - POST into remote CPA Management API (`cpa_auto_import_remote`)
"""

from __future__ import annotations

import os
import shutil
import sys
import time
from pathlib import Path
from typing import Any, Callable
from urllib.parse import quote, urlparse, urlunparse

_REG_DIR = Path(__file__).resolve().parent
_DEFAULT_OUT = _REG_DIR / "cpa_auths"
_DEFAULT_CPA = Path("")  # empty = do not assume a machine-local CPA path


def _ensure_cpa_xai_on_path(tools_dir: str | Path | None = None) -> Path:
    """Put the parent of `cpa_xai` on sys.path. Default: this project root."""
    if tools_dir:
        tools = Path(tools_dir).expanduser().resolve()
    else:
        env = (os.environ.get("API_REVERSE_TOOLS") or "").strip()
        tools = Path(env).expanduser().resolve() if env else _REG_DIR
    # If user pointed at .../cpa_xai itself, use its parent
    if tools.name == "cpa_xai" and (tools / "__init__.py").is_file():
        tools = tools.parent
    if str(tools) not in sys.path:
        sys.path.insert(0, str(tools))
    return tools


def export_cookies_from_page(page: Any) -> list[dict]:
    """Best-effort export of cookies from a DrissionPage tab/browser."""
    if page is None:
        return []
    cookies = None
    for getter in (
        lambda: page.cookies(all_domains=True, all_info=True),
        lambda: page.cookies(all_domains=True),
        lambda: page.cookies(),
    ):
        try:
            cookies = getter()
            if cookies:
                break
        except TypeError:
            continue
        except Exception:
            continue
    if not cookies:
        try:
            browser = getattr(page, "browser", None)
            if browser is not None:
                cookies = browser.cookies()
        except Exception:
            cookies = None
    if isinstance(cookies, list):
        return [c for c in cookies if isinstance(c, dict)]
    return []


def _normalize_cpa_remote_base(base: str) -> str:
    """Normalize user-provided CPA root URL to scheme://host[:port].

    Accepts:
      http://127.0.0.1:8317
      http://host:8317/
      http://host:8317/v0/management
      http://host:8317/v0/management/auth-files
    """
    raw = (base or "").strip()
    if not raw:
        return ""
    if "://" not in raw:
        raw = "http://" + raw
    parsed = urlparse(raw)
    if not parsed.scheme or not parsed.netloc:
        return raw.rstrip("/")
    # Keep only scheme + netloc (drop path/query/fragment)
    return urlunparse((parsed.scheme, parsed.netloc, "", "", "", "")).rstrip("/")


def resolve_cpa_remote_settings(cfg: dict | None = None) -> dict[str, Any]:
    """Resolve remote CPA import settings from config + env."""
    cfg = cfg or {}
    enabled = bool(cfg.get("cpa_auto_import_remote", False))
    base = (
        (cfg.get("cpa_remote_base") or "").strip()
        or (os.environ.get("CPA_REMOTE_BASE") or "").strip()
    )
    password = (
        (cfg.get("cpa_remote_password") or "").strip()
        or (os.environ.get("CPA_REMOTE_PASSWORD") or "").strip()
        or (os.environ.get("CPA_MANAGEMENT_PASSWORD") or "").strip()
        or (os.environ.get("MANAGEMENT_PASSWORD") or "").strip()
    )
    retries = int(cfg.get("cpa_remote_import_retries", 3) or 3)
    delay = float(cfg.get("cpa_remote_import_retry_delay", 2) or 2)
    return {
        "enabled": enabled,
        "base": _normalize_cpa_remote_base(base),
        "password": password,
        "retries": max(1, retries),
        "delay": max(0.0, delay),
    }


def import_cpa_auth_to_remote(
    path: str | Path,
    *,
    base: str,
    password: str,
    retries: int = 3,
    delay: float = 2.0,
    log: Callable[[str], None] | None = None,
    timeout: float = 12.0,
) -> dict[str, Any]:
    """Upload one local xai-*.json to CPA Management API.

    POST {base}/v0/management/auth-files?name=<filename>
    Authorization: Bearer <password>
    Body: raw JSON file content

    Failures are returned in the result dict; caller decides whether to ignore.
    """
    _log = log or (lambda m: print(m, flush=True))
    src = Path(path)
    root = _normalize_cpa_remote_base(base)
    if not root:
        return {"ok": False, "error": "empty cpa_remote_base"}
    if not password:
        return {"ok": False, "error": "empty cpa_remote_password"}
    if not src.is_file():
        return {"ok": False, "error": f"auth file not found: {src}"}

    name = src.name
    if not name.endswith(".json"):
        return {"ok": False, "error": f"filename must end with .json: {name}"}

    try:
        body = src.read_bytes()
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": f"read file: {e}"}

    url = f"{root}/v0/management/auth-files?name={quote(name)}"
    headers = {
        "Authorization": f"Bearer {password}",
        "X-Management-Key": password,
        "Content-Type": "application/json",
    }

    try:
        from curl_cffi import requests as _req  # type: ignore
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": f"import curl_cffi failed: {e}"}

    last_err = "unknown"
    attempts = max(1, int(retries or 1))
    for i in range(1, attempts + 1):
        try:
            # Admin API is local/remote management — do not use outbound proxy.
            resp = _req.post(
                url,
                data=body,
                headers=headers,
                timeout=timeout,
                proxies={},
                impersonate="chrome",
            )
            status = int(getattr(resp, "status_code", 0) or 0)
            text = ""
            try:
                text = (resp.text or "")[:500]
            except Exception:
                text = ""
            if status in (200, 207):
                # Prefer JSON status field when present
                ok_status = True
                try:
                    j = resp.json()
                    st = str((j or {}).get("status") or "").lower()
                    if st and st not in ("ok", "partial", "success"):
                        ok_status = False
                        last_err = f"status={st} body={text}"
                except Exception:
                    pass
                if ok_status:
                    _log(f"[cpa] remote import ok -> {root} name={name}")
                    return {
                        "ok": True,
                        "base": root,
                        "name": name,
                        "status_code": status,
                        "body": text,
                    }
            else:
                last_err = f"HTTP {status}: {text or '(empty)'}"
        except Exception as e:  # noqa: BLE001
            last_err = str(e)

        if i < attempts:
            _log(f"[cpa] remote import retry {i}/{attempts}: {last_err}")
            if delay > 0:
                time.sleep(delay)

    _log(f"[cpa] remote import failed: {last_err}")
    return {
        "ok": False,
        "base": root,
        "name": name,
        "error": last_err,
    }


def maybe_import_cpa_auth_remote(
    path: str | Path,
    *,
    config: dict | None = None,
    log_callback: Callable[[str], None] | None = None,
) -> dict[str, Any]:
    """Config-gated remote import helper (safe no-op when disabled)."""
    cfg = config or {}
    log = log_callback or (lambda m: print(m, flush=True))
    settings = resolve_cpa_remote_settings(cfg)
    if not settings["enabled"]:
        return {"ok": False, "skipped": True, "reason": "disabled"}
    if not settings["base"] or not settings["password"]:
        log(
            "[cpa] remote import enabled but missing cpa_remote_base / "
            "cpa_remote_password (or CPA_REMOTE_BASE / CPA_REMOTE_PASSWORD)"
        )
        return {
            "ok": False,
            "skipped": True,
            "reason": "missing_base_or_password",
        }
    return import_cpa_auth_to_remote(
        path,
        base=settings["base"],
        password=settings["password"],
        retries=settings["retries"],
        delay=settings["delay"],
        log=log,
    )


def export_cpa_xai_for_account(
    email: str,
    password: str,
    *,
    page: Any | None = None,
    cookies: Any | None = None,
    sso: str | None = None,
    config: dict | None = None,
    log_callback: Callable[[str], None] | None = None,
) -> dict:
    """Mint OIDC + write xai-<email>.json under register cpa_auths (and optional CPA auth-dir)."""
    cfg = config or {}
    log = log_callback or (lambda m: print(m, flush=True))

    if not cfg.get("cpa_export_enabled", True):
        log("[cpa] export disabled")
        return {"ok": False, "skipped": True, "reason": "disabled"}

    tools_dir = cfg.get("api_reverse_tools") or cfg.get("cpa_xai_parent") or None
    _ensure_cpa_xai_on_path(tools_dir)

    try:
        from cpa_xai import mint_and_export  # type: ignore
    except Exception as e:  # noqa: BLE001
        log(f"[cpa] import cpa_xai failed: {e}")
        return {"ok": False, "error": f"import: {e}"}

    out_dir = Path(cfg.get("cpa_auth_dir") or _DEFAULT_OUT).expanduser()
    if not out_dir.is_absolute():
        out_dir = (_REG_DIR / out_dir).resolve()

    hotload_raw = (cfg.get("cpa_hotload_dir") or "").strip()
    cpa_dir = Path(hotload_raw).expanduser() if hotload_raw else None
    if cpa_dir and not cpa_dir.is_absolute():
        cpa_dir = (_REG_DIR / cpa_dir).resolve()

    # Priority: cpa_proxy > proxy > env. Config must beat shell https_proxy.
    proxy = (cfg.get("cpa_proxy") or cfg.get("proxy") or "").strip()
    if not proxy:
        proxy = (
            os.environ.get("https_proxy")
            or os.environ.get("HTTPS_PROXY")
            or os.environ.get("http_proxy")
            or ""
        ).strip()
    # Default headed: headless is frequently Cloudflare-blocked on accounts.x.ai
    headless = bool(cfg.get("cpa_headless", False))
    probe = bool(cfg.get("cpa_probe_after_write", True))
    probe_chat = bool(cfg.get("cpa_probe_chat", False))
    timeout = float(cfg.get("cpa_mint_timeout_sec", 240))
    base_url = cfg.get("cpa_base_url") or "https://cli-chat-proxy.grok.com/v1"
    force_standalone = bool(cfg.get("cpa_force_standalone", True))
    cookie_inject = bool(cfg.get("cpa_mint_cookie_inject", True))
    reuse_browser = bool(cfg.get("cpa_mint_browser_reuse", True))
    recycle_every = int(cfg.get("cpa_mint_browser_recycle_every", 15) or 0)
    # Protocol (pure HTTP SSO device flow) first; browser only on failure.
    prefer_protocol = bool(cfg.get("cpa_prefer_protocol", True))
    protocol_only = bool(cfg.get("cpa_protocol_only", False))
    protocol_poll_timeout = float(cfg.get("cpa_protocol_poll_timeout_sec", 90) or 90)

    # cookies: explicit arg > page export > none
    use_cookies = cookies
    if use_cookies is None and cookie_inject and page is not None:
        use_cookies = export_cookies_from_page(page)
    if not cookie_inject:
        use_cookies = None
    else:
        # Always attach SSO cookie clones — register cookies alone often miss accounts.x.ai host
        sso_val = (sso or "").strip()
        if not sso_val and isinstance(use_cookies, list):
            for c in use_cookies:
                if isinstance(c, dict) and c.get("name") in ("sso", "sso-rw") and c.get("value"):
                    sso_val = str(c.get("value"))
                    break
        if sso_val:
            base = list(use_cookies) if isinstance(use_cookies, list) else []
            for name in ("sso", "sso-rw"):
                for dom in (".x.ai", "accounts.x.ai", ".accounts.x.ai", "auth.x.ai", "grok.com", ".grok.com"):
                    base.append({
                        "name": name,
                        "value": sso_val,
                        "domain": dom,
                        "path": "/",
                        "secure": True,
                        "httpOnly": True,
                    })
            use_cookies = base

    sso_val = (sso or "").strip()
    if not sso_val and isinstance(use_cookies, list):
        for c in use_cookies:
            if isinstance(c, dict) and c.get("name") in ("sso", "sso-rw") and c.get("value"):
                sso_val = str(c.get("value"))
                break

    out_dir.mkdir(parents=True, exist_ok=True)
    log(
        f"[cpa] mint OIDC for {email} -> {out_dir} proxy={proxy or '(none)'} "
        f"cookies={len(use_cookies) if isinstance(use_cookies, list) else (1 if use_cookies else 0)} "
        f"reuse={reuse_browser} protocol={prefer_protocol}"
        f"{' only' if protocol_only else ''} sso={'yes' if sso_val else 'no'}"
    )

    def _log(msg: str) -> None:
        log(f"[cpa] {msg}")

    result = mint_and_export(
        email=email,
        password=password,
        auth_dir=out_dir,
        page=None if force_standalone else page,
        proxy=proxy or None,
        headless=headless,
        base_url=base_url,
        probe=probe,
        probe_chat=probe_chat,
        browser_timeout_sec=timeout,
        force_standalone=force_standalone,
        cookies=use_cookies,
        sso=sso_val or None,
        reuse_browser=reuse_browser,
        recycle_every=recycle_every,
        prefer_protocol=prefer_protocol,
        protocol_only=protocol_only,
        protocol_poll_timeout_sec=protocol_poll_timeout,
        log=_log,
    )
    if result.get("mint_method"):
        log(f"[cpa] mint_method={result.get('mint_method')}")

    # By default, a failed post-write probe is only a warning: the CPA auth file
    # has already been minted and written. Set cpa_probe_required=true to make
    # missing /models grok-4.5 fail the export.
    if (
        not result.get("ok")
        and result.get("path")
        and str(result.get("error") or "").startswith("token ok but grok-4.5 not listed")
        and not cfg.get("cpa_probe_required", False)
    ):
        result["ok"] = True
        result["probe_warning"] = result.pop("error", "probe failed")
        log(f"[cpa] probe warning ignored (file already written): {result.get('probe_warning')}")

    if result.get("ok") and result.get("path") and cfg.get("cpa_copy_to_hotload", False) and cpa_dir:
        try:
            cpa_dir.mkdir(parents=True, exist_ok=True)
            src = Path(result["path"])
            dst = cpa_dir / src.name
            shutil.copy2(src, dst)
            os.chmod(dst, 0o600)
            result["cpa_path"] = str(dst)
            log(f"[cpa] hotload copy -> {dst}")
        except Exception as e:  # noqa: BLE001
            log(f"[cpa] hotload copy failed: {e}")
            result["cpa_copy_error"] = str(e)

    # Optional: push auth file into remote CPA via Management API.
    # Failures never flip result["ok"] — local mint already succeeded.
    if result.get("ok") and result.get("path"):
        try:
            remote = maybe_import_cpa_auth_remote(
                result["path"],
                config=cfg,
                log_callback=log,
            )
            if remote.get("skipped"):
                result["remote_import_skipped"] = remote.get("reason") or "skipped"
            elif remote.get("ok"):
                result["remote_import_ok"] = True
                result["remote_import_base"] = remote.get("base")
                result["remote_import_name"] = remote.get("name")
            else:
                result["remote_import_ok"] = False
                result["remote_import_error"] = remote.get("error") or "unknown"
        except Exception as e:  # noqa: BLE001
            log(f"[cpa] remote import unexpected error: {e}")
            result["remote_import_ok"] = False
            result["remote_import_error"] = str(e)

    # failure log under register dir
    if not result.get("ok"):
        fail_path = out_dir / "cpa_auth_failed.txt"
        with open(fail_path, "a", encoding="utf-8") as f:
            f.write(f"{email}----{result.get('error') or 'unknown'}----{int(time.time())}\n")
        if cfg.get("cpa_mint_required", False):
            raise RuntimeError(f"CPA mint required but failed: {result.get('error')}")

    return result