#!/usr/bin/env python3
"""Replace local Turnstile/browser CF path with pure zip implementations.

Keeps local Gmail / hotmail / register_cli / cpa_xai intact.
Source of truth: grok-register.zip (ReinerBRO) for CF + browser options.
"""
from __future__ import annotations

import os
import re
import shutil
from pathlib import Path

ROOT = Path(r"D:/Projects/AI/grok-keygen")
# Git Bash /tmp maps to %TEMP% on Windows
_TEMP = Path(os.environ.get("TEMP") or os.environ.get("TMP") or r"C:/Users/Chaos/AppData/Local/Temp")
ZIP_REG = _TEMP / "zip_full/grok-register/grok_register_ttk.py"
ZIP_CPA = _TEMP / "zip_full/grok_reg-protocol_cpa/grok_register_ttk.py"
ZIP_PATCH = _TEMP / "zip_full/grok-register/turnstilePatch"
LOCAL = ROOT / "grok_register_ttk.py"
LOCAL_PATCH = ROOT / "turnstilePatch"


def extract_func(text: str, name: str) -> str:
    lines = text.splitlines(keepends=True)
    starts = {}
    for i, line in enumerate(lines):
        m = re.match(r"^def (\w+)\(", line)
        if m:
            starts[m.group(1)] = i
        m = re.match(r"^([A-Z_][A-Z0-9_]*)\s*=", line)
        if m:
            starts[m.group(1)] = i
    if name not in starts:
        raise KeyError(f"missing {name}")
    s = starts[name]
    e = len(lines)
    for j in range(s + 1, len(lines)):
        if re.match(r"^(def |class |[A-Z_][A-Z0-9_]*\s*=)", lines[j]):
            e = j
            break
    return "".join(lines[s:e]).rstrip() + "\n\n\n"


def replace_func(src: str, name: str, new_body: str) -> str:
    """Replace a top-level def/assignment block by name."""
    lines = src.splitlines(keepends=True)
    start = None
    for i, line in enumerate(lines):
        if re.match(rf"^def {re.escape(name)}\(", line) or re.match(
            rf"^{re.escape(name)}\s*=", line
        ):
            start = i
            break
    if start is None:
        raise KeyError(f"local missing {name}")
    end = len(lines)
    for j in range(start + 1, len(lines)):
        if re.match(r"^(def |class |[A-Z_][A-Z0-9_]*\s*=)", lines[j]):
            end = j
            break
    # preserve trailing blank style: ensure exactly two newlines before next
    body = new_body
    if not body.endswith("\n"):
        body += "\n"
    # normalize to end with two blank lines if next block exists
    while body.endswith("\n\n\n"):
        body = body[:-1]
    if not body.endswith("\n\n"):
        body = body.rstrip("\n") + "\n\n"
    return "".join(lines[:start]) + body + "".join(lines[end:])


def remove_func(src: str, name: str) -> str:
    lines = src.splitlines(keepends=True)
    start = None
    for i, line in enumerate(lines):
        if re.match(rf"^def {re.escape(name)}\(", line) or re.match(
            rf"^{re.escape(name)}\s*=", line
        ):
            start = i
            break
    if start is None:
        return src
    end = len(lines)
    for j in range(start + 1, len(lines)):
        if re.match(r"^(def |class |[A-Z_][A-Z0-9_]*\s*=)", lines[j]):
            end = j
            break
    return "".join(lines[:start]) + "".join(lines[end:])


def adapt_reg_get_turnstile(zip_fn: str) -> str:
    """grok-register uses session.page; local uses _get_page()."""
    text = zip_fn
    text = text.replace(
        "def getTurnstileToken(session, log_callback=None, cancel_callback=None):\n    page = session.page\n",
        "def getTurnstileToken(log_callback=None, cancel_callback=None, max_rounds: int = 20):\n    page = _get_page()\n",
    )
    # support max_rounds
    text = text.replace(
        "for _ in range(0, 20):",
        "rounds = max(1, int(max_rounds or 20))\n    for _ in range(0, rounds):",
    )
    # local uses human_sleep/cf_sleep; zip uses sleep_with_cancel — keep sleep_with_cancel (exists)
    return text


def adapt_reg_fill_profile(zip_fn: str) -> str:
    """Adapt fill_profile from grok-register to local globals."""
    text = zip_fn
    text = text.replace(
        "def fill_profile_and_submit(session, timeout=120, log_callback=None, cancel_callback=None):\n    page = session.page\n",
        "def fill_profile_and_submit(timeout=120, log_callback=None, cancel_callback=None):\n    page = _get_page()\n    check_timeout(time.time())\n    dump_state(page, \"profile-form\")\n    take_screenshot(page, \"profile-form\")\n",
    )
    # fix getTurnstileToken(session, ...) -> getTurnstileToken(...)
    text = text.replace(
        "token = getTurnstileToken(session, log_callback=log_callback, cancel_callback=cancel_callback)",
        "token = getTurnstileToken(log_callback=log_callback, cancel_callback=cancel_callback)",
    )
    # add zip warmup like protocol_cpa
    if "预热 Turnstile" not in text:
        text = text.replace(
            "given_name, family_name, password = build_profile()\n    deadline = time.time() + timeout\n",
            "given_name, family_name, password = build_profile()\n"
            "    if log_callback:\n"
            "        log_callback(\"[*] 预热 Turnstile...\")\n"
            "    human_sleep(2, cancel_callback)\n"
            "    deadline = time.time() + timeout\n",
        )
    return text


def adapt_create_browser_options(zip_fn: str) -> str:
    """Use zip minimal options; keep local headless/proxy helpers lightly."""
    # Prefer ReinerBRO pure options (no disable-images). Keep proxy + headless glue.
    return '''def create_browser_options(headless=None):
    """Zip-style Chromium options (grok-register / ReinerBRO).

    Critical for CF:
      - always load turnstilePatch
      - no --disable-images
      - no browser_path unless config/env set
      - no AutomationControlled / heavy stealth flags on headed
    """
    options = ChromiumOptions()
    options.auto_port()
    options.set_timeouts(base=1)
    # ReinerBRO headed flags
    for flag in (
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-background-networking",
        "--disable-component-update",
        "--disable-background-timer-throttling",
        "--disable-backgrounding-occluded-windows",
        "--disable-renderer-backgrounding",
    ):
        options.set_argument(flag)

    use_headless = should_use_headless(headless, log_hint=True)
    if use_headless:
        apply_headless_to_options(options, True)
        for flag in (
            "--no-sandbox",
            "--disable-dev-shm-usage",
            "--disable-gpu",
            "--mute-audio",
        ):
            options.set_argument(flag)
        print("  [browser] headless=true", flush=True)
    else:
        print("  [browser] headless=false (zip headed)", flush=True)

    # system default unless override
    browser_path = resolve_browser_path()
    if browser_path:
        try:
            options.set_browser_path(browser_path)
            print(f"  [browser] path={browser_path}", flush=True)
        except Exception as exc:
            print(f"  [browser] set path failed: {exc}", flush=True)
    else:
        print("  [browser] path=system-default (zip style)", flush=True)

    # Force UA from config (Chrome/138 etc.)
    try:
        ua = get_user_agent()
        if ua:
            options.set_user_agent(ua)
            print(f"  [browser] ua={ua[:60]}...", flush=True)
    except Exception:
        pass

    if os.path.exists(EXTENSION_PATH):
        try:
            options.add_extension(EXTENSION_PATH)
            print(f"  [ext] turnstilePatch loaded: {EXTENSION_PATH}", flush=True)
        except Exception as ext_exc:
            print(f"  [ext] turnstilePatch load failed: {ext_exc}", flush=True)
    else:
        print(f"  [ext] turnstilePatch missing: {EXTENSION_PATH}", flush=True)

    proxy = (config.get("proxy") or "").strip() if isinstance(config, dict) else ""
    if proxy:
        try:
            from urllib.parse import urlparse

            u = urlparse(proxy if "://" in proxy else f"http://{proxy}")
            host = u.hostname or ""
            if host:
                port = u.port or (443 if (u.scheme or "http") == "https" else 80)
                scheme = u.scheme or "http"
                options.set_argument(f"--proxy-server={scheme}://{host}:{port}")
                print(f"  [proxy] browser --proxy-server={scheme}://{host}:{port}", flush=True)
        except Exception as e:
            print(f"  [proxy] set browser proxy failed: {e}", flush=True)
    return options


'''


def main() -> None:
    if not ZIP_REG.is_file():
        raise SystemExit(f"missing {ZIP_REG}")
    reg_text = ZIP_REG.read_text(encoding="utf-8", errors="replace")
    local = LOCAL.read_text(encoding="utf-8", errors="replace")

    # backup
    bak = ROOT / "grok_register_ttk.py.bak_pre_zip_port"
    shutil.copy2(LOCAL, bak)
    print(f"backup -> {bak}")

    # restore pure zip turnstilePatch
    if ZIP_PATCH.is_dir():
        LOCAL_PATCH.mkdir(exist_ok=True)
        for name in ("content.js", "manifest.json"):
            src = ZIP_PATCH / name
            if src.is_file():
                shutil.copy2(src, LOCAL_PATCH / name)
                print(f"copied {name}")
        # remove MV3 leftovers
        for stale in ("script.js",):
            p = LOCAL_PATCH / stale
            if p.exists():
                p.unlink()
                print(f"removed {stale}")

    get_ts = adapt_reg_get_turnstile(extract_func(reg_text, "getTurnstileToken"))
    fill = adapt_reg_fill_profile(extract_func(reg_text, "fill_profile_and_submit"))
    create_opts = adapt_create_browser_options("")

    # Remove interactive-safe helpers that zip doesn't need (keep if referenced elsewhere carefully)
    for helper in (
        "_click_turnstile_widget",
        "_read_turnstile_token_simple",
        "_turnstile_ui_state",
        "_read_turnstile_token",
        "_turnstile_diag",
    ):
        before = local
        local = remove_func(local, helper)
        if local != before:
            print(f"removed helper {helper}")

    local = replace_func(local, "getTurnstileToken", get_ts)
    print("replaced getTurnstileToken")
    local = replace_func(local, "fill_profile_and_submit", fill)
    print("replaced fill_profile_and_submit")
    local = replace_func(local, "create_browser_options", create_opts)
    print("replaced create_browser_options")

    # Simplify resolve_browser_path to zip style (auto -> None)
    resolve_browser = '''def resolve_browser_path():
    """Zip style: only return path if config/env override exists.

    browser_prefer=auto/system → None (DrissionPage system Chrome).
    """
    override = str((config.get("browser_path") if isinstance(config, dict) else "") or "").strip()
    if override:
        norm = _normalize_browser_path(override)
        if norm:
            return norm
        print(f"  [browser] ignore invalid browser_path={override!r}", flush=True)

    for env_key in ("BROWSER_PATH", "CHROME_PATH", "CHROMIUM_PATH"):
        env_path = (os.environ.get(env_key) or "").strip()
        if not env_path:
            continue
        norm = _normalize_browser_path(env_path)
        if norm:
            return norm
        print(f"  [browser] ignore invalid {env_key}={env_path!r}", flush=True)

    prefer = str((config.get("browser_prefer") if isinstance(config, dict) else "") or "auto").strip().lower()
    if prefer in ("", "default", "detect", "system"):
        prefer = "auto"
    env_pref = (os.environ.get("BROWSER_PREFER") or "").strip().lower()
    if env_pref:
        prefer = env_pref

    if prefer in ("auto", "system"):
        return None
    if prefer in ("chrome", "google-chrome", "google"):
        order = list(CHROME_CANDIDATES) + list(CHROMIUM_CANDIDATES)
    elif prefer in ("chromium", "chromeium"):
        order = list(CHROMIUM_CANDIDATES) + list(CHROME_CANDIDATES)
    else:
        return None
    return _first_existing_browser(order)


'''
    local = replace_func(local, "resolve_browser_path", resolve_browser)
    print("replaced resolve_browser_path")

    # Force stealth off by default in apply_stealth_patches
    local = re.sub(
        r'mode = str\(\(config\.get\("stealth_mode"\).*?or "\w+"\)\.strip\(\)\.lower\(\)',
        'mode = str((config.get("stealth_mode") if isinstance(config, dict) else "") or "off").strip().lower()',
        local,
        count=1,
    )

    # Ensure fill_profile / other call sites don't reference removed helpers
    # getTurnstileToken(..., max_rounds=) still ok
    if "_turnstile_ui_state" in local or "_click_turnstile_widget" in local:
        print("WARN: leftover helper references:")
        for i, line in enumerate(local.splitlines(), 1):
            if "_turnstile_ui_state" in line or "_click_turnstile_widget" in line or "_read_turnstile_token" in line:
                print(f"  L{i}: {line.strip()[:120]}")

    LOCAL.write_text(local, encoding="utf-8")
    print(f"wrote {LOCAL} ({len(local.splitlines())} lines)")

    # syntax check
    import py_compile

    py_compile.compile(str(LOCAL), doraise=True)
    print("syntax OK")


if __name__ == "__main__":
    main()