#!/usr/bin/env python3
"""
Minimal CDP-over-pipe client.  Launches Chrome with --remote-debugging-pipe,
navigates to a URL, waits for load, captures a screenshot, and exits.

Usage:
  python3 cdp_screenshot.py <chrome-binary> <url> <out.png> [extra chrome args...]
"""
import base64
import json
import os
import subprocess
import sys
import threading
import time

CHROME = sys.argv[1]
URL = sys.argv[2]
OUT = sys.argv[3]
EXTRA = sys.argv[4:]

# fd 3: chrome reads from here (we write).  fd 4: chrome writes to here (we read).
r_to_chrome, w_to_chrome = os.pipe()
r_from_chrome, w_from_chrome = os.pipe()

cmd = [
    CHROME,
    "--headless",
    "--no-sandbox",
    "--disable-gpu",
    "--no-first-run",
    "--disable-crash-reporter",
    "--remote-debugging-pipe",
    f"--user-data-dir=/tmp/cdp-prof-{os.getpid()}",
    "about:blank",
] + EXTRA

env = dict(os.environ)
env["ASAN_OPTIONS"] = "detect_leaks=0"
env.setdefault("HOME", "/tmp/chrome-home")

os.set_inheritable(r_to_chrome, True)
os.set_inheritable(w_from_chrome, True)


def _child_setup():
    os.dup2(r_to_chrome, 3)
    os.dup2(w_from_chrome, 4)
    # leave originals open; harmless


proc = subprocess.Popen(
    cmd,
    stdin=subprocess.DEVNULL,
    stdout=subprocess.DEVNULL,
    stderr=open("/tmp/cdp-stderr.log", "wb"),
    close_fds=False,
    preexec_fn=_child_setup,
)
os.close(r_to_chrome)
os.close(w_from_chrome)

write_f = os.fdopen(w_to_chrome, "wb")
read_f = os.fdopen(r_from_chrome, "rb")

_id = [0]
_pending = {}
_events = []
_lock = threading.Lock()
_cv = threading.Condition(_lock)


def reader():
    buf = b""
    while True:
        chunk = read_f.read1(65536)
        if not chunk:
            break
        buf += chunk
        while b"\x00" in buf:
            msg, buf = buf.split(b"\x00", 1)
            if not msg:
                continue
            obj = json.loads(msg)
            with _cv:
                if "id" in obj:
                    _pending[obj["id"]] = obj
                else:
                    _events.append(obj)
                _cv.notify_all()


threading.Thread(target=reader, daemon=True).start()


def send(method, params=None, session_id=None, timeout=30):
    _id[0] += 1
    mid = _id[0]
    msg = {"id": mid, "method": method, "params": params or {}}
    if session_id:
        msg["sessionId"] = session_id
    write_f.write(json.dumps(msg).encode() + b"\x00")
    write_f.flush()
    deadline = time.time() + timeout
    with _cv:
        while mid not in _pending:
            if not _cv.wait(timeout=deadline - time.time()):
                raise TimeoutError(f"{method} timed out")
        return _pending.pop(mid)


def wait_event(method, session_id=None, timeout=30):
    deadline = time.time() + timeout
    with _cv:
        while True:
            for i, e in enumerate(_events):
                if e.get("method") == method and (
                    session_id is None or e.get("sessionId") == session_id
                ):
                    return _events.pop(i)
            if not _cv.wait(timeout=max(0, deadline - time.time())):
                raise TimeoutError(f"event {method} timed out")


try:
    r = send("Target.createTarget", {"url": "about:blank"})
    target_id = r["result"]["targetId"]
    r = send("Target.attachToTarget", {"targetId": target_id, "flatten": True})
    sid = r["result"]["sessionId"]

    send("Page.enable", session_id=sid)
    send(
        "Emulation.setDeviceMetricsOverride",
        {"width": 400, "height": 400, "deviceScaleFactor": 1, "mobile": False},
        session_id=sid,
    )
    send("Page.navigate", {"url": URL}, session_id=sid)
    wait_event("Page.loadEventFired", session_id=sid, timeout=30)
    time.sleep(0.5)  # let paint settle

    r = send(
        "Page.captureScreenshot",
        {"format": "png", "clip": {"x": 0, "y": 0, "width": 400, "height": 400, "scale": 1}},
        session_id=sid,
        timeout=30,
    )
    data = base64.b64decode(r["result"]["data"])
    with open(OUT, "wb") as f:
        f.write(data)
    print(f"wrote {OUT} ({len(data)} bytes)")
finally:
    try:
        send("Browser.close", timeout=5)
    except Exception:
        pass
    proc.terminate()
    try:
        proc.wait(timeout=5)
    except Exception:
        proc.kill()