#!/usr/bin/env python3
"""Reproduce output-removal crashes only in a private headless compositor."""
import argparse
import json
import os
from pathlib import Path
import resource
import socket
import struct
import subprocess
import tempfile
import time

parser = argparse.ArgumentParser()
parser.add_argument('--mode', choices=['window', 'native', 'forced'], required=True)
parser.add_argument('--plugins', type=Path, required=True)
parser.add_argument('--gdb', action='store_true')
parser.add_argument('--cycle', action='store_true')
parser.add_argument('--order', default='213')
parser.add_argument('--return-order', default='123')
parser.add_argument('--clients', type=int, default=1)
parser.add_argument('--cycles', type=int, default=1)
parser.add_argument('--pause', type=float, default=.1)
parser.add_argument('--inhibit', action='store_true')
parser.add_argument('--unload', action='store_true')
args = parser.parse_args()
assert sorted(args.order) == sorted(args.return_order) == list('123')
assert os.geteuid() != 0, 'Run as an ordinary user'
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
runtime = Path(tempfile.mkdtemp(prefix='wayfire-crash-repro-'))
env = os.environ.copy()
for key in ['WAYLAND_DISPLAY', 'DISPLAY', 'WAYFIRE_SOCKET', '_WAYFIRE_SOCKET',
            'WAYFIRE_CONFIG_FILE', 'WAYFIRE_DEFAULT_CONFIG_BACKEND', 'LD_LIBRARY_PATH']:
    env.pop(key, None)
assert env.get('DBUS_SESSION_BUS_ADDRESS'), 'Use dbus-run-session'
env.update(XDG_RUNTIME_DIR=str(runtime), XDG_CONFIG_HOME=str(runtime / 'config'),
           DBUS_SYSTEM_BUS_ADDRESS=env['DBUS_SESSION_BUS_ADDRESS'],
           WLR_BACKENDS='headless', WLR_HEADLESS_OUTPUTS='3', WLR_RENDERER='gles2',
           WLR_RENDER_DRM_DEVICE='/dev/dri/renderD128',
           WAYFIRE_PLUGIN_PATH=str(args.plugins), GDK_BACKEND='wayland')
config = runtime / 'wayfire.ini'
config.write_text('''[core]
plugins = ipc ipc-rules window-rules stipc wm-actions force-fullscreen hotplug-test
xwayland = false
vwidth = 1
vheight = 1
stitch_enabled = true
stitch_outputs = HEADLESS-1 HEADLESS-2 HEADLESS-3
default_output = VIRTUAL-1
[force-fullscreen]
key_toggle_fullscreen = <alt> <super> KEY_F
constrain_pointer = false
preserve_aspect = true
transparent_behind_views = true
[output:HEADLESS-1]
mode = 640x480@60
position = 0,0
[output:HEADLESS-2]
mode = 640x480@60
position = 640,0
[output:HEADLESS-3]
mode = 640x480@60
position = 1280,0
''')

def ipc(method, data=None):
    path = env['WAYFIRE_SOCKET']
    assert Path(path).parent == runtime
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
        s.settimeout(3)
        s.connect(path)
        payload = json.dumps({'method': method, 'data': data or {}}).encode()
        s.sendall(struct.pack('<I', len(payload)) + payload)
        def read(n):
            buf = b''
            while len(buf) < n:
                part = s.recv(n - len(buf))
                if not part:
                    raise EOFError('Private compositor disconnected')
                buf += part
            return buf
        result = json.loads(read(struct.unpack('<I', read(4))[0]))
        assert not isinstance(result, dict) or not result.get('error'), result
        return result

def wait(fn, timeout=10):
    end = time.monotonic() + timeout
    while time.monotonic() < end:
        assert process.poll() is None, f'Exited early: {runtime}'
        value = fn()
        if value:
            return value
        time.sleep(.05)
    raise AssertionError(f'Timed out: {runtime}')

command = ['/usr/bin/wayfire', '-c', str(config)]
if args.gdb:
    command = ['gdb', '-q', '-batch', '-ex', 'set pagination off', '-ex', 'set confirm off',
               '-ex', 'run', '-ex', 'bt 24', '--args', *command]
clients = []
print(runtime, flush=True)
with (runtime / 'wayfire.log').open('w') as log:
    process = subprocess.Popen(command, env=env, cwd=runtime, stdout=log, stderr=log)
    try:
        def connect():
            sockets = list(runtime.glob('wayfire-*.socket'))
            displays = [p for p in runtime.glob('wayland-*') if not p.name.endswith('.lock')]
            if sockets and displays:
                env['WAYFIRE_SOCKET'] = str(sockets[0])
                env['WAYLAND_DISPLAY'] = displays[0].name
                return True
        wait(connect)
        wait(lambda: any(o['name'] == 'VIRTUAL-1' for o in ipc('window-rules/list-outputs')))
        with (runtime / 'client.log').open('w') as client_log:
            for _ in range(args.clients):
                clients.append(subprocess.Popen(['alacritty', '--class', 'crash-test', '-e', '/bin/sleep', '300'],
                                               env=env, cwd=runtime, stdout=client_log, stderr=client_log))
        def views():
            return [v for v in ipc('window-rules/list-views') if v.get('app-id') == 'crash-test']
        wait(lambda: len(vs := views()) == args.clients and all(v['mapped'] for v in vs))
        ids = {v['id'] for v in views()}
        def set_fullscreen():
            for vid in sorted(ids):
                ipc('window-rules/focus-view', {'id': vid})
                wait(lambda: any(v['id'] == vid and v['activated'] for v in views()))
                if args.mode == 'native':
                    ipc('wm-actions/set-fullscreen', {'view_id': vid, 'state': True})
                elif args.mode == 'forced':
                    for key in ['KEY_LEFTALT', 'KEY_LEFTMETA', 'KEY_F']:
                        ipc('stipc/feed_key', {'key': key, 'state': True})
                    for key in ['KEY_F', 'KEY_LEFTMETA', 'KEY_LEFTALT']:
                        ipc('stipc/feed_key', {'key': key, 'state': False})
                if args.mode != 'window':
                    wait(lambda: any(v['id'] == vid and v['fullscreen'] for v in views()))
            if args.mode == 'forced' and args.cycle:
                assert {v['id'] for v in ipc('hotplug-test/state') if v['forced']} == ids
        set_fullscreen()
        time.sleep(.4)
        (runtime / 'before.json').write_text(json.dumps(ipc('window-rules/list-views'), indent=2))
        if args.unload:
            original_config = config.read_text()
            if args.inhibit:
                ipc('hotplug-test/inhibit', {'inhibited': True})
            config.write_text(original_config.replace('wm-actions force-fullscreen', 'wm-actions'))
            wait(lambda: not any(v['forced'] for v in ipc('hotplug-test/state')))
            wait(lambda: all(not v['fullscreen'] for v in views()))
            assert {v['id'] for v in views()} == ids
            ipc('hotplug-test/inhibit', {'inhibited': False})
            config.write_text(original_config)
            time.sleep(.5)
            set_fullscreen()
            result = {'mode': args.mode, 'passed': True, 'unload': True, 'inhibited': args.inhibit,
                      'runtime': str(runtime)}
            (runtime / 'result.json').write_text(json.dumps(result, indent=2))
            print('PASS plugin unload/reload and fullscreen reuse', flush=True)
            raise SystemExit(0)
        if args.cycle:
            states = []
            for cycle in range(args.cycles):
                if cycle:
                    set_fullscreen()
                if args.inhibit:
                    ipc('hotplug-test/inhibit', {'inhibited': True})
                for number in args.order:
                    ipc('hotplug-test/remove', {'name': 'HEADLESS-' + number})
                    time.sleep(args.pause)
                wait(lambda: all(o['name'].startswith('NOOP-') for o in ipc('window-rules/list-outputs')))
                off = views()
                assert {v['id'] for v in off} == ids and all(p.poll() is None for p in clients), off
                assert not any(v['forced'] for v in ipc('hotplug-test/state'))
                for number in args.return_order:
                    ipc('hotplug-test/add', {'name': 'HEADLESS-' + number})
                    time.sleep(args.pause)
                wait(lambda: [o['name'] for o in ipc('window-rules/list-outputs')] == ['VIRTUAL-1'])
                restored = wait(lambda: (vs := views()) and all(v['mapped'] and v['output-name'] == 'VIRTUAL-1' for v in vs) and vs)
                assert {v['id'] for v in restored} == ids and all(p.poll() is None for p in clients)
                states.append({'cycle': cycle, 'off': off, 'restored': restored})
                print(f'PASS cycle {cycle + 1}: off {args.order}, on {args.return_order}, {args.clients} {args.mode} window(s)', flush=True)
            result = {'mode': args.mode, 'passed': True, 'states': states, 'runtime': str(runtime)}
            (runtime / 'result.json').write_text(json.dumps(result, indent=2))
            raise SystemExit(0)
        try:
            ipc('hotplug-test/remove', {'name': 'HEADLESS-2'})
        except (EOFError, ConnectionError, TimeoutError):
            pass
        try:
            code = process.wait(timeout=3)
            result = {'mode': args.mode, 'exited': True, 'code': code, 'runtime': str(runtime)}
        except subprocess.TimeoutExpired:
            result = {'mode': args.mode, 'exited': False, 'outputs': ipc('window-rules/list-outputs'),
                      'runtime': str(runtime)}
        (runtime / 'result.json').write_text(json.dumps(result, indent=2))
        print(json.dumps(result), flush=True)
    finally:
        for client in clients:
            if client.poll() is None:
                client.terminate()
                client.wait(timeout=5)
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()