#!/usr/bin/env python3
"""Headless chromium CDP driver: load URL, capture console, wait for __sceneReady, screenshot.
usage: shoot.py URL OUT.png [--w 1280] [--h 720] [--wait MS] [--eval JS] [--frames N] [--timeout S] [--mobile]
"""
import asyncio, json, subprocess, sys, base64, time, urllib.request, re, argparse, os, signal
import websockets

ap = argparse.ArgumentParser()
ap.add_argument('url'); ap.add_argument('out')
ap.add_argument('--w', type=int, default=1280); ap.add_argument('--h', type=int, default=720)
ap.add_argument('--wait', type=int, default=0, help='extra ms of real time after ready before shot')
ap.add_argument('--eval', default=None, help='JS to evaluate after ready (before extra wait)')
ap.add_argument('--timeout', type=float, default=240)
ap.add_argument('--mobile', action='store_true')
ap.add_argument('--dpr', type=float, default=1.0)
ap.add_argument('--virtual', action='store_true', help='deterministic virtual clock: 1/30s per frame, frozen after ready')
args = ap.parse_args()

flags = ['/usr/bin/chromium', '--headless=new', '--no-sandbox', '--disable-dev-shm-usage',
         '--use-angle=swiftshader', '--enable-unsafe-swiftshader', '--ignore-gpu-blocklist',
         '--remote-debugging-port=0', '--hide-scrollbars', f'--window-size={args.w},{args.h}',
         '--no-first-run', '--disable-extensions', '--mute-audio', '--user-data-dir=/tmp/chrome-prof-%d' % os.getpid(), 'about:blank']
proc = subprocess.Popen(flags, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
ws_url = None
t0 = time.time()
while time.time() - t0 < 30:
    line = proc.stderr.readline()
    if not line: time.sleep(0.05); continue
    m = re.search(r'DevTools listening on (ws://\S+)', line)
    if m: ws_url = m.group(1); break
if not ws_url:
    print('no devtools url'); proc.kill(); sys.exit(1)
port = re.search(r':(\d+)/', ws_url).group(1)

async def main():
    targets = json.loads(urllib.request.urlopen(f'http://127.0.0.1:{port}/json').read())
    page = [t for t in targets if t['type'] == 'page'][0]
    msgs = []
    async with websockets.connect(page['webSocketDebuggerUrl'], max_size=2**30) as ws:
        mid = [0]
        pending = {}
        async def send(method, params=None):
            mid[0] += 1
            i = mid[0]
            fut = asyncio.get_event_loop().create_future(); pending[i] = fut
            await ws.send(json.dumps({'id': i, 'method': method, 'params': params or {}}))
            return await fut
        async def reader():
            async for raw in ws:
                m = json.loads(raw)
                if 'id' in m and m['id'] in pending:
                    pending.pop(m['id']).set_result(m.get('result', m.get('error')))
                elif m.get('method') == 'Runtime.consoleAPICalled':
                    p = m['params']
                    txt = ' '.join(str(a.get('value', a.get('description', a.get('type')))) for a in p['args'])
                    msgs.append(f"[console.{p['type']}] {txt[:600]}")
                elif m.get('method') == 'Runtime.exceptionThrown':
                    d = m['params']['exceptionDetails']
                    msgs.append(f"[EXCEPTION] {d.get('text')} {(d.get('exception') or {}).get('description','')[:800]}")
                elif m.get('method') == 'Log.entryAdded':
                    e = m['params']['entry']
                    msgs.append(f"[log.{e['level']}] {e.get('source')}: {e['text'][:600]} {e.get('url','')}")
        rt = asyncio.create_task(reader())
        await send('Runtime.enable'); await send('Log.enable'); await send('Page.enable')
        await send('Emulation.setDeviceMetricsOverride', {'width': args.w, 'height': args.h, 'deviceScaleFactor': args.dpr, 'mobile': args.mobile})
        if args.mobile:
            await send('Emulation.setTouchEmulationEnabled', {'enabled': True, 'maxTouchPoints': 5})
        if args.virtual:
            await send('Page.addScriptToEvaluateOnNewDocument', {'source': '''(function(){
              let now = 0; window.__vstep = 1000/30;
              performance.now = function(){ return now; };
              const _raf = window.requestAnimationFrame.bind(window);
              window.requestAnimationFrame = function(cb){ return _raf(function(){ now += window.__vstep; cb(now); }); };
            })();'''})
        await send('Page.navigate', {'url': args.url})
        t1 = time.time(); ready = False
        while time.time() - t1 < args.timeout:
            r = await send('Runtime.evaluate', {'expression': 'window.__sceneReady === true', 'returnByValue': True})
            if r and r.get('result', {}).get('value') is True: ready = True; break
            await asyncio.sleep(0.5)
        print('ready:', ready, 'after %.1fs' % (time.time() - t1))
        if args.virtual:
            await send('Runtime.evaluate', {'expression': 'window.__vstep = 0; true', 'returnByValue': True})
            await asyncio.sleep(1.5)
        if args.eval:
            r = await send('Runtime.evaluate', {'expression': args.eval, 'returnByValue': True, 'awaitPromise': True})
            print('eval:', json.dumps(r)[:1000])
        if args.wait: await asyncio.sleep(args.wait / 1000)
        r = await send('Page.captureScreenshot', {'format': 'png'})
        with open(args.out, 'wb') as f: f.write(base64.b64decode(r['data']))
        print('saved', args.out)
        stats = await send('Runtime.evaluate', {'expression': 'JSON.stringify(window.__stats||null)', 'returnByValue': True})
        print('stats:', (stats or {}).get('result', {}).get('value'))
        rt.cancel()
    for m_ in msgs: print(m_)
    print(f'-- {len(msgs)} console/log messages')

try:
    asyncio.run(main())
finally:
    proc.kill()
    try: proc.wait(timeout=5)
    except Exception: pass
    subprocess.run(['rm', '-rf', '/tmp/chrome-prof-%d' % os.getpid()])