"""Drive a headless LibreWolf that carries the owner's live OpenAI session.

The console pages gptbridge must reach (platform tunnels, ChatGPT developer apps)
sit behind Cloudflare and a React SPA, so a cookie jar replayed through curl gets a
403 challenge instead of a session. This snapshots the owner's profile, runs a
second browser against the copy, and speaks Marionette to it.
"""

from __future__ import annotations

import fcntl
import json
import os
import secrets
import select
import shutil
import signal
import socket
import subprocess
import tempfile
import time
from pathlib import Path

HOME = Path.home()
STATE = Path(os.environ.get("GPTBRIDGE_STATE_DIR", HOME / ".overdeck" / "gptbridge"))
PROFILE_COPY = STATE / "browser-profile"
DOWNLOAD_DIR = STATE / "downloads"

# Caches and lock files never carry session state and dominate the copy cost.
SKIP = {
    "cache2",
    "startupCache",
    "shader-cache",
    "thumbnails",
    "crashes",
    "minidumps",
    "datareporting",
    "saved-telemetry-pings",
    "lock",
    ".parentlock",
}


class MarionetteError(RuntimeError):
    pass


def source_profile() -> Path:
    """The profile the owner's running browser uses, per profiles.ini."""
    ini = HOME / ".librewolf/profiles.ini"
    if not ini.is_file():
        raise MarionetteError(f"no LibreWolf profiles.ini at {ini}")
    default, section_path, sections = None, None, {}
    name = None
    for raw in ini.read_text().splitlines():
        line = raw.strip()
        if line.startswith("["):
            if name is not None:
                sections[name] = section_path
            name, section_path = line, None
        elif line.startswith("Path="):
            section_path = line[5:]
        elif line.startswith("Default=") and name and name.startswith("[Install"):
            default = line[8:]
    if name is not None:
        sections[name] = section_path
    chosen = default or next((p for p in sections.values() if p), None)
    if not chosen:
        raise MarionetteError("profiles.ini names no profile")
    path = HOME / ".librewolf" / chosen
    if not path.is_dir():
        raise MarionetteError(f"profile {chosen} missing at {path}")
    return path


# Applied on every run, including a reused copy, so a pref change takes effect
# without discarding the profile's verification cookies.
AUTOMATION_PREFS = """\
user_pref("marionette.port", %(port)d);
user_pref("browser.shell.checkDefaultBrowser", false);
user_pref("browser.sessionstore.resume_from_crash", false);
user_pref("datareporting.policy.firstRunURL", "");
// Marionette otherwise sets navigator.webdriver, which the bot check on the OpenAI
// consoles treats as an unsolvable challenge.
user_pref("dom.webdriver.enabled", false);
// LibreWolf pins the content area to a rounded size and letterboxes it; a pointer
// click below that fold is then rejected as out of viewport. Throwaway copy only —
// never the owner's own profile.
user_pref("privacy.resistFingerprinting", false);
user_pref("privacy.resistFingerprinting.letterboxing", false);
// No window manager runs on the virtual display, so the window cannot be resized and
// a tall menu falls outside it. Shrinking the CSS pixel fits more page in the same
// window instead.
user_pref("layout.css.devPixelsPerPx", "0.7");
user_pref("browser.download.folderList", 2);
user_pref("browser.download.dir", %(download_dir)s);
user_pref("browser.download.useDownloadDir", true);
user_pref("browser.download.alwaysOpenPanel", false);
user_pref("browser.helperApps.neverAsk.saveToDisk", "application/zip,application/octet-stream");
"""


XVFB_SCREEN = ["-screen", "0", "1600x1200x24", "-nolisten", "tcp"]


def _orphan(cmdline: bytes, profile: Path) -> bool:
    if str(profile).encode() in cmdline and b"librewolf" in cmdline.lower():
        return True
    parts = cmdline.split(b"\0")
    return bool(parts) and parts[0].endswith(b"Xvfb") and all(
        arg.encode() in parts for arg in XVFB_SCREEN)


def reap_orphans(profile: Path) -> int:
    """Terminate browsers and virtual displays left by a run that was killed.

    Only ever called while this process holds the profile lock, so a live run cannot
    be the target; matching is on the profile path and this module's own Xvfb
    signature, so nothing else the owner runs can be either.
    """
    killed = 0
    for entry in Path("/proc").iterdir():
        if not entry.name.isdigit():
            continue
        try:
            cmdline = (entry / "cmdline").read_bytes()
        except OSError:
            continue
        if not _orphan(cmdline, profile):
            continue
        try:
            os.kill(int(entry.name), signal.SIGTERM)
            killed += 1
        except OSError:
            continue
    if killed:
        time.sleep(4)
    return killed


def free_port() -> int:
    for _ in range(256):
        port = secrets.randbelow(65535 - 49152 + 1) + 49152
        with socket.socket() as sock:
            try:
                sock.bind(("127.0.0.1", port))
            except OSError:
                continue
        return port
    raise MarionetteError("no private dynamic port is available")


def snapshot_profile(src: Path, dest: Path = PROFILE_COPY, reseed: bool = False,
                     port: int = 2828) -> Path:
    """Copy the live profile so the owner's running browser is never touched.

    Kept across runs: reseeding discards the automation profile's own verification
    cookies, which makes every run face the bot check again.
    """
    if reseed and dest.exists():
        shutil.rmtree(dest)
    if not dest.is_dir():
        dest.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        shutil.copytree(
            src,
            dest,
            ignore=lambda _d, names: [n for n in names if n in SKIP],
            symlinks=True,
            ignore_dangling_symlinks=True,
        )
        os.chmod(dest, 0o700)
    DOWNLOAD_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
    (dest / "user.js").write_text(AUTOMATION_PREFS % {
        "port": port,
        "download_dir": json.dumps(str(DOWNLOAD_DIR)),
    })
    return dest


class Marionette:
    """Minimal Marionette client: length-prefixed JSON over TCP."""

    def __init__(self, port: int = 2828, host: str = "127.0.0.1") -> None:
        self.host, self.port = host, port
        self.sock: socket.socket | None = None
        self.buf = b""
        self.msgid = 0

    def connect(self, timeout: float = 60.0) -> None:
        deadline = time.monotonic() + timeout
        last: Exception | None = None
        while time.monotonic() < deadline:
            try:
                self.sock = socket.create_connection((self.host, self.port), 5)
                self.sock.settimeout(180)
                self._read_packet()  # server handshake
                self.command("WebDriver:NewSession", {})
                # Pointer clicks are rejected outside the viewport, so the window must
                # not be left at whatever size the window manager hands out.
                self.command(
                    "WebDriver:SetWindowRect",
                    {"x": 0, "y": 0, "width": 1400, "height": 1000},
                )
                return
            except OSError as exc:
                last = exc
                self.sock = None
                time.sleep(0.5)
        raise MarionetteError(f"marionette never came up on {self.port}: {last}")

    def _read_packet(self) -> object:
        assert self.sock is not None
        while b":" not in self.buf:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise MarionetteError("marionette closed the connection")
            self.buf += chunk
        size, _, rest = self.buf.partition(b":")
        need = int(size)
        while len(rest) < need:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise MarionetteError("marionette closed mid-packet")
            rest += chunk
        self.buf = rest[need:]
        return json.loads(rest[:need])

    def command(self, name: str, params: dict) -> object:
        if self.sock is None:
            raise MarionetteError("marionette session is closed")
        self.msgid += 1
        body = json.dumps([0, self.msgid, name, params]).encode()
        self.sock.sendall(b"%d:%s" % (len(body), body))
        while True:
            packet = self._read_packet()
            if not isinstance(packet, list) or len(packet) != 4 or packet[0] != 1:
                continue
            _, mid, err, result = packet
            if mid != self.msgid:
                continue
            if err:
                raise MarionetteError(json.dumps(err))
            return result.get("value") if isinstance(result, dict) else result

    def navigate(self, url: str) -> None:
        self.command("WebDriver:Navigate", {"url": url})

    def url(self) -> str:
        return str(self.command("WebDriver:GetCurrentURL", {}))

    def script(self, source: str, args: list | None = None) -> object:
        return self.command(
            "WebDriver:ExecuteAsyncScript",
            {"script": source, "args": args or [], "newSandbox": False},
        )

    def sync_script(self, source: str, args: list | None = None) -> object:
        return self.command(
            "WebDriver:ExecuteScript",
            {"script": source, "args": args or [], "newSandbox": False},
        )

    def find(self, selector: str) -> str | None:
        try:
            el = self.command(
                "WebDriver:FindElement", {"using": "css selector", "value": selector}
            )
        except MarionetteError:
            return None
        if isinstance(el, dict):
            for k, v in el.items():
                if "element" in k:
                    return str(v)
        return None

    def click(self, selector: str) -> bool:
        """Click an element, falling back to a pointer click at its centre.

        Menu surfaces built on floating-overlay libraries report as not scrollable
        into view, which fails WebDriver's element click but not a real pointer.
        """
        ref = self.find(selector)
        if ref is None:
            return False
        try:
            self.command("WebDriver:ElementClick", {"id": ref})
            return True
        except MarionetteError as exc:
            if "not interactable" not in str(exc):
                raise
        rect = self.sync_script(
            "const e=document.querySelector(arguments[0]);"
            "if(!e)return null;const r=e.getBoundingClientRect();"
            "if(r.width<1||r.height<1)return null;"
            "return {x:r.x+r.width/2,y:r.y+r.height/2};",
            [selector],
        )
        if not isinstance(rect, dict):
            return False
        self.click_point(rect["x"], rect["y"])
        return True

    JS_CLICK = """
const e = document.querySelector(arguments[0]);
if (!e) return false;
const r = e.getBoundingClientRect();
const base = {bubbles: true, cancelable: true, view: window, button: 0,
              clientX: r.x + r.width / 2, clientY: r.y + r.height / 2};
for (const type of ['pointerover', 'pointerenter', 'pointermove', 'pointerdown',
                    'mousedown', 'pointerup', 'mouseup', 'click']) {
  const Ctor = type.startsWith('pointer') ? PointerEvent : MouseEvent;
  const init = type.startsWith('pointer')
    ? {...base, pointerType: 'mouse', isPrimary: true} : base;
  e.dispatchEvent(new Ctor(type, init));
}
return true;
"""

    def js_click(self, selector: str) -> bool:
        """Click by dispatching the full pointer sequence in the page.

        No window manager runs on the virtual display, so the window cannot be
        resized and a tall menu extends past the viewport, where WebDriver refuses
        to move the pointer. Dispatching in-page has no such bound.
        """
        return bool(self.sync_script(self.JS_CLICK, [selector]))

    ESCAPE = "\ue00c"

    def press(self, key: str = ESCAPE) -> None:
        self.command("WebDriver:PerformActions", {"actions": [{
            "type": "key", "id": "kb",
            "actions": [{"type": "keyDown", "value": key},
                        {"type": "keyUp", "value": key}]}]})
        self.command("WebDriver:ReleaseActions", {})

    def send_keys(self, selector: str, text: str) -> bool:
        ref = self.find(selector)
        if ref is None:
            return False
        self.command("WebDriver:ElementSendKeys", {"id": ref, "text": text})
        return True

    def click_point(self, x: float, y: float) -> None:
        """Click at viewport coordinates through the input layer.

        Needed for widgets inside a cross-origin iframe, which no in-page script can
        reach.
        """
        self.command(
            "WebDriver:PerformActions",
            {
                "actions": [
                    {
                        "type": "pointer",
                        "id": "mouse",
                        "parameters": {"pointerType": "mouse"},
                        "actions": [
                            {"type": "pointerMove", "duration": 300,
                             "origin": "viewport", "x": int(x), "y": int(y)},
                            {"type": "pause", "duration": 150},
                            {"type": "pointerDown", "button": 0},
                            {"type": "pause", "duration": 90},
                            {"type": "pointerUp", "button": 0},
                        ],
                    }
                ]
            },
        )
        self.command("WebDriver:ReleaseActions", {})

    def screenshot(self, dest: Path) -> Path:
        import base64

        data = self.command("WebDriver:TakeScreenshot", {"full": True, "hash": False})
        dest.write_bytes(base64.b64decode(str(data)))
        return dest

    def close(self) -> None:
        sock, self.sock = self.sock, None
        if sock is not None:
            try:
                self.sock = sock
                self.command("Marionette:Quit", {})
            except (MarionetteError, OSError):
                pass
            finally:
                self.sock = None
            sock.close()


def start_xvfb(timeout: float = 20.0) -> tuple[subprocess.Popen, int]:
    """Start a virtual display and return it with the number Xvfb chose.

    Xvfb picks the number itself and reports it on -displayfd. Choosing one here
    from the socket paths under /tmp/.X11-unix is not sound: a caller's /tmp may be
    a sandbox overlay that hides another display's socket, while the abstract socket
    Xvfb binds is shared with every namespace — an unseen server then makes the new
    one exit at once with "server already running".
    """
    read_fd, write_fd = os.pipe()
    errors = tempfile.TemporaryFile()
    try:
        proc = subprocess.Popen(
            ["Xvfb", "-displayfd", str(write_fd), *XVFB_SCREEN],
            pass_fds=(write_fd,), stdout=subprocess.DEVNULL, stderr=errors,
        )
    finally:
        os.close(write_fd)

    reported = b""
    deadline = time.monotonic() + timeout
    while b"\n" not in reported and time.monotonic() < deadline:
        if select.select([read_fd], [], [], 0.2)[0]:
            chunk = os.read(read_fd, 64)
            if not chunk:
                break
            reported += chunk
            continue
        if proc.poll() is not None:
            break
    os.close(read_fd)

    display = reported.strip().decode(errors="replace")
    if not display.isdigit():
        proc.terminate()
        try:
            proc.wait(5)
        except subprocess.TimeoutExpired:
            proc.kill()
        errors.seek(0)
        detail = errors.read().decode(errors="replace").strip().splitlines()
        raise MarionetteError(
            "Xvfb never reported a display"
            + (f": {detail[-1]}" if detail else " and printed no error")
        )
    return proc, int(display)


class Session:
    """Run the copied profile in one of three display modes.

    ``virtual`` is the default: Cloudflare's interstitial never clears for a
    ``--headless`` Gecko, but the same build on an Xvfb display passes and stays
    off the owner's screen.
    """

    MODES = ("virtual", "headless", "show")
    LOCK_WAIT_SECONDS = 600

    def __init__(self, port: int | None = None, mode: str = "virtual",
                 reseed: bool = False) -> None:
        if mode not in self.MODES:
            raise MarionetteError(f"mode must be one of {self.MODES}")
        self.port = port or free_port()
        self.mode, self.reseed = mode, reseed
        self.lock = None
        self.proc: subprocess.Popen | None = None
        self.xvfb: subprocess.Popen | None = None
        self.client = Marionette(self.port)

    def __enter__(self) -> Marionette:
        STATE.mkdir(mode=0o700, parents=True, exist_ok=True)
        os.chmod(STATE, 0o700)
        # One profile directory, so only one browser may use it at a time.
        self.lock = (STATE / "profile.lock").open("w")
        deadline = time.monotonic() + self.LOCK_WAIT_SECONDS
        while True:
            try:
                fcntl.flock(self.lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except OSError:
                if time.monotonic() > deadline:
                    raise MarionetteError(
                        "another ask-gpt run is using the browser profile"
                    ) from None
                time.sleep(2)
        # Holding the lock proves no live run owns the profile, so anything still
        # attached to it is an orphan from a run that was killed.
        reap_orphans(PROFILE_COPY)
        profile = snapshot_profile(source_profile(), reseed=self.reseed,
                                   port=self.port)
        env = dict(os.environ)
        env.pop("MOZ_CRASHREPORTER", None)
        # Without a window manager the window is whatever Gecko picks, and a pointer
        # click outside the viewport is rejected outright.
        argv = ["librewolf", "--profile", str(profile), "--no-remote", "--marionette",
                "-width", "1500", "-height", "1100"]

        if self.mode == "headless":
            env["MOZ_HEADLESS"] = "1"
            argv += ["--headless", "--window-size=1500,1100"]
        elif self.mode == "virtual":
            env.pop("MOZ_HEADLESS", None)
            self.xvfb, display = start_xvfb()
            env["DISPLAY"] = f":{display}"
        else:
            env.pop("MOZ_HEADLESS", None)

        self.proc = subprocess.Popen(
            argv, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        self.client.connect()
        return self.client

    def __exit__(self, *_exc: object) -> None:
        self.client.close()
        if self.proc is not None:
            try:
                self.proc.wait(15)
            except subprocess.TimeoutExpired:
                self.proc.kill()
                self.proc.wait(10)
        if self.xvfb is not None:
            self.xvfb.terminate()
            try:
                self.xvfb.wait(10)
            except subprocess.TimeoutExpired:
                self.xvfb.kill()
        if self.lock is not None:
            self.lock.close()
            self.lock = None