"""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 select
import shutil
import signal
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path

HOME = Path.home()
STATE = HOME / ".overdeck" / "gptbridge"
PROFILE_COPY = STATE / "browser-profile"
DOWNLOAD_DIR = STATE / "downloads"

# Slot 1 keeps the original unsuffixed paths, so the default (one seat) is
# byte-identical to before parallel seats existed. Slots 2+ get their own copy,
# lock and download dir. Each open slot costs a LibreWolf + Xvfb (~0.5-1 GB RSS)
# and its first use pays a profile copy; slots only exist once actually opened.
MAX_SLOTS = 15


def profile_path(slot: int) -> Path:
    return PROFILE_COPY if slot == 1 else STATE / f"browser-profile-{slot}"


def lock_path(slot: int) -> Path:
    return STATE / "profile.lock" if slot == 1 else STATE / f"profile-{slot}.lock"


def download_path(slot: int) -> Path:
    return DOWNLOAD_DIR if slot == 1 else STATE / f"downloads-{slot}"


def _xvfb_pid_marker(slot: int) -> Path:
    return STATE / f"xvfb-{slot}.pid"

# 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");
user_pref("dom.events.testing.asyncClipboard", true);
user_pref("permissions.default.clipboard-read", 1);
"""


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


def _orphan_browser(cmdline: bytes, profile: Path) -> bool:
    # A raw substring test on the whole cmdline blob would match slot 1's unsuffixed
    # path against slot 2's ("browser-profile" is a substring of "browser-profile-2"),
    # reaping a live sibling seat. Match a whole argv element, or a path under it,
    # instead.
    if b"librewolf" not in cmdline.lower():
        return False
    want = str(profile).encode()
    return any(part == want or part.startswith(want + b"/")
              for part in cmdline.split(b"\0"))


def _orphan_xvfb_pid(slot: int) -> int | None:
    """The PID this slot's own last Xvfb was recorded under, if it still looks live.

    Xvfb's argv carries no display number with -displayfd, and its screen-mode
    signature is identical across every slot, so matching by signature alone would
    also catch a live sibling seat's display. The PID this slot itself wrote is the
    only thing that names one slot's Xvfb without naming another's.
    """
    marker = _xvfb_pid_marker(slot)
    if not marker.is_file():
        return None
    try:
        pid = int(marker.read_text().strip())
        parts = (Path("/proc") / str(pid) / "cmdline").read_bytes().split(b"\0")
    except (OSError, ValueError):
        marker.unlink(missing_ok=True)
        return None
    if not (parts and parts[0].endswith(b"Xvfb")):
        marker.unlink(missing_ok=True)
        return None
    return pid


def reap_orphans(profile: Path, slot: int = 1) -> int:
    """Terminate this slot's browser and virtual display left by a run that was killed.

    Only ever called while this process holds the slot's lock, so a live run of THIS
    slot cannot be the target. Browser matching is scoped by the slot's own profile
    path; Xvfb matching is scoped to the PID this slot itself recorded, never a live
    sibling seat's.
    """
    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_browser(cmdline, profile):
            continue
        try:
            os.kill(int(entry.name), signal.SIGTERM)
            killed += 1
        except OSError:
            continue
    pid = _orphan_xvfb_pid(slot)
    if pid is not None:
        try:
            os.kill(pid, signal.SIGTERM)
            killed += 1
        except OSError:
            pass
        _xvfb_pid_marker(slot).unlink(missing_ok=True)
    if killed:
        time.sleep(4)
    return killed


def free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return int(s.getsockname()[1])


def snapshot_profile(src: Path, dest: Path = PROFILE_COPY, reseed: bool = False,
                     port: int = 2828, download_dir: Path = DOWNLOAD_DIR) -> 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
        # Overwritten by Session.__enter__ with the owning slot's own download dir;
        # this default only matters for a Marionette built outside a Session.
        self.download_dir = DOWNLOAD_DIR

    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 two 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. Headless is not an offered mode for that reason —
    a caller asking for it silently gets ``virtual`` instead.

    ``show`` puts a window on the owner's screen, so it runs only when the owner
    sets ``OVERDECK_BROWSER_SHOW=1``; every other caller is coerced to
    ``virtual``.
    """

    MODES = ("virtual", "show")
    LOCK_WAIT_SECONDS = 600
    SHOW_ENV = "OVERDECK_BROWSER_SHOW"

    def __init__(self, port: int | None = None, mode: str = "virtual",
                 reseed: bool = False, slot: int | None = None) -> None:
        if mode not in self.MODES:
            raise MarionetteError(f"mode must be one of {self.MODES}")
        if mode == "show" and os.environ.get(self.SHOW_ENV) != "1":
            print(f"notice: --mode show needs {self.SHOW_ENV}=1; using virtual",
                  file=sys.stderr)
            mode = "virtual"
        # A slot's own lock serialises every browser that could touch it, so a port
        # picked fresh per session is never contended.
        self.port = port or free_port()
        self.mode, self.reseed = mode, reseed
        # None = self-pick the first free slot on __enter__; an explicit slot (the
        # pool's per-seat Sessions) is pinned and never probed.
        self.slot = slot
        self.lock = None
        self.proc: subprocess.Popen | None = None
        self.xvfb: subprocess.Popen | None = None
        self.client = Marionette(self.port)

    @staticmethod
    def _flock_wait(handle, deadline: float) -> None:
        while True:
            try:
                fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
                return
            except OSError:
                if time.monotonic() > deadline:
                    raise MarionetteError(
                        "another ask-gpt run is using the browser profile"
                    ) from None
                time.sleep(2)

    def _acquire_slot(self) -> int:
        if self.slot is not None:
            self.lock = lock_path(self.slot).open("w")
            self._flock_wait(self.lock, time.monotonic() + self.LOCK_WAIT_SECONDS)
            return self.slot
        # Non-blocking probe across every slot first: only wait if all are taken.
        for slot in range(1, MAX_SLOTS + 1):
            handle = lock_path(slot).open("w")
            try:
                fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except OSError:
                handle.close()
                continue
            self.lock = handle
            return slot
        self.lock = lock_path(1).open("w")
        self._flock_wait(self.lock, time.monotonic() + self.LOCK_WAIT_SECONDS)
        return 1

    def __enter__(self) -> Marionette:
        STATE.mkdir(mode=0o700, parents=True, exist_ok=True)
        self.slot = self._acquire_slot()
        # Holding the lock proves no live run owns this slot, so anything still
        # attached to it is an orphan from a run that was killed.
        reap_orphans(profile_path(self.slot), self.slot)
        profile = snapshot_profile(source_profile(), dest=profile_path(self.slot),
                                   reseed=self.reseed, port=self.port,
                                   download_dir=download_path(self.slot))
        self.client.download_dir = download_path(self.slot)
        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 == "virtual":
            env.pop("MOZ_HEADLESS", None)
            self.xvfb, display = start_xvfb()
            env["DISPLAY"] = f":{display}"
            # Read back by reap_orphans if this run is killed before __exit__ runs.
            _xvfb_pid_marker(self.slot).write_text(str(self.xvfb.pid))
        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()
            _xvfb_pid_marker(self.slot).unlink(missing_ok=True)
        if self.lock is not None:
            self.lock.close()
            self.lock = None