import base64, json, pathlib, subprocess, sys, tempfile, time, urllib.request
import websocket
profile, output, start, end = sys.argv[1:]
with tempfile.TemporaryDirectory(prefix='ipu-profile-browser-') as directory:
    log = open(pathlib.Path(output).with_suffix('.browser.log'), 'w')
    process = subprocess.Popen(['chromium', '--headless', '--no-sandbox', '--disable-gpu', '--disable-extensions', '--disable-background-networking', '--remote-debugging-port=0', '--user-data-dir='+directory, 'about:blank'], stdout=log, stderr=log)
    try:
        port_file = pathlib.Path(directory)/'DevToolsActivePort'
        deadline=time.monotonic()+30
        while not port_file.exists():
            if time.monotonic()>deadline: raise TimeoutError('browser startup')
            time.sleep(.1)
        port=port_file.read_text().splitlines()[0]
        tabs=json.load(urllib.request.urlopen('http://127.0.0.1:'+port+'/json'))
        ws=websocket.create_connection(next(t for t in tabs if t['type']=='page')['webSocketDebuggerUrl'], suppress_origin=True, timeout=60)
        serial=0
        def call(method, params={}):
            global serial
            serial+=1
            ws.send(json.dumps(dict(id=serial,method=method,params=params)))
            while True:
                message=json.loads(ws.recv())
                if message.get('id')==serial:
                    if 'error' in message: raise RuntimeError(message['error'])
                    return message.get('result',{})
        call('Emulation.setDeviceMetricsOverride', dict(width=2000,height=2000,deviceScaleFactor=1,mobile=False))
        call('Page.navigate', dict(url=pathlib.Path(profile).resolve().as_uri()))
        deadline=time.monotonic()+60
        while True:
            result=call('Runtime.evaluate',dict(expression='typeof applyView',returnByValue=True))
            if result.get('result',{}).get('value')=='function':break
            if time.monotonic()>deadline:raise TimeoutError('profile loading')
            time.sleep(.2)
        result=call('Runtime.evaluate',dict(expression=f'applyView({int(start)}, {int(end)}); document.getElementById("timeline").toDataURL("image/png")',returnByValue=True))
        data=result['result']['value'].split(',',1)[1]
        pathlib.Path(output).write_bytes(base64.b64decode(data))
        print(output)
        ws.close()
    finally:
        process.terminate()
        process.wait(timeout=10)
        log.close()