#!/usr/bin/env python3
"""
Drive Chrome via --remote-debugging-pipe (CDP over fd 3/4, NUL-delimited JSON)
to load the POC page and write a screenshot to ./screenshot.png.

No external deps (no websocket libs needed).
"""
import base64
import json
import os
import subprocess
import sys
import time

CHROME = os.environ.get(
    'CHROME', '/chromium/chromium/src/out/asan/chrome')
URL = sys.argv[1] if len(sys.argv) > 1 else 'http://127.0.0.1:39080/poc.html'
OUT = sys.argv[2] if len(sys.argv) > 2 else os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 'screenshot.png')

import fcntl
# fd 3: chrome reads (we write); fd 4: chrome writes (we read).
# Allocate pipes on high fds so we can place chrome's ends on 3/4 without
# clobbering the parent-side ends.
def move_above(fd, lo=10):
    nfd = fcntl.fcntl(fd, fcntl.F_DUPFD, lo)
    os.close(fd)
    return nfd

r3, w3 = os.pipe()   # we write w3, chrome reads r3 (as its fd 3)
r4, w4 = os.pipe()   # chrome writes w4 (as its fd 4), we read r4
r3 = move_above(r3); w3 = move_above(w3)
r4 = move_above(r4); w4 = move_above(w4)

argv = [
    CHROME,
    '--headless', '--no-sandbox', '--disable-gpu',
    '--no-first-run', '--no-default-browser-check',
    '--user-data-dir=/tmp/chrome-cdp-profile',
    '--force-device-scale-factor=1',
    '--window-size=500,350',
    '--hide-scrollbars',
    '--remote-debugging-pipe',
    '--remote-allow-origins=*',
    'about:blank',
]
print('launching:', ' '.join(argv), file=sys.stderr)
errlog = open('/tmp/chrome-cdp-stderr.log', 'wb')

# Put the pipe ends on fds 3 and 4 in *this* process so the child inherits
# them positionally (Chrome checks fd 3/4 very early, before preexec hooks
# can be relied on across Python versions).
os.dup2(r3, 3); os.close(r3)
os.dup2(w4, 4); os.close(w4)
os.set_inheritable(3, True)
os.set_inheritable(4, True)

proc = subprocess.Popen(
    argv,
    pass_fds=(3, 4),
    stdin=subprocess.DEVNULL,
    stdout=subprocess.DEVNULL,
    stderr=errlog,
)
# Parent no longer needs the chrome-side ends.
os.close(3)
os.close(4)

buf = b''
def read_msg(timeout=60):
    global buf
    import select
    end = time.time() + timeout
    while True:
        nl = buf.find(b'\0')
        if nl >= 0:
            msg, buf = buf[:nl], buf[nl+1:]
            return json.loads(msg)
        rem = end - time.time()
        if rem <= 0:
            raise TimeoutError('CDP read timeout')
        r, _, _ = select.select([r4], [], [], rem)
        if not r:
            raise TimeoutError('CDP read timeout')
        chunk = os.read(r4, 65536)
        if not chunk:
            raise EOFError('CDP pipe closed')
        buf += chunk

mid = 0
def send(method, params=None, session=None):
    global mid
    mid += 1
    m = {'id': mid, 'method': method}
    if params is not None:
        m['params'] = params
    if session is not None:
        m['sessionId'] = session
    os.write(w3, json.dumps(m).encode() + b'\0')
    return mid

def wait_id(want, timeout=60):
    end = time.time() + timeout
    while time.time() < end:
        m = read_msg(end - time.time())
        if m.get('id') == want:
            return m
    raise TimeoutError(f'no reply for id {want}')

def wait_event(name, session=None, timeout=60):
    end = time.time() + timeout
    while time.time() < end:
        m = read_msg(end - time.time())
        if m.get('method') == name and (session is None or m.get('sessionId') == session):
            return m
    raise TimeoutError(f'no event {name}')

try:
    # Create a page target.
    i = send('Target.createTarget', {'url': 'about:blank', 'width': 500, 'height': 350})
    r = wait_id(i)
    target_id = r['result']['targetId']

    i = send('Target.attachToTarget', {'targetId': target_id, 'flatten': True})
    r = wait_id(i)
    sess = r['result']['sessionId']

    # Force a 500x350 layout viewport so absolute box coordinates match.
    wait_id(send('Emulation.setDeviceMetricsOverride', {
        'width': 500, 'height': 350, 'deviceScaleFactor': 1, 'mobile': False
    }, sess))

    # Enable page events and navigate.
    wait_id(send('Page.enable', {}, sess))
    wait_id(send('Page.navigate', {'url': URL}, sess))
    wait_event('Page.loadEventFired', sess, timeout=60)

    # Give background images time to decode/paint, then force a frame.
    time.sleep(2)
    wait_id(send('Runtime.evaluate', {
        'expression': 'new Promise(r=>requestAnimationFrame(()=>requestAnimationFrame(r)))',
        'awaitPromise': True}, sess), timeout=30)

    # Capture.
    i = send('Page.captureScreenshot', {
        'format': 'png', 'fromSurface': True,
        'clip': {'x': 0, 'y': 0, 'width': 500, 'height': 350, 'scale': 1}
    }, sess)
    r = wait_id(i, timeout=60)
    data = base64.b64decode(r['result']['data'])
    with open(OUT, 'wb') as f:
        f.write(data)
    print(f'wrote {len(data)} bytes to {OUT}', file=sys.stderr)

    send('Browser.close')
    time.sleep(1)
finally:
    try:
        proc.terminate()
        proc.wait(timeout=10)
    except Exception:
        proc.kill()