#!/usr/bin/env python3
"""Two disposable KVM guests on a private headless display; never host input."""
import argparse
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile
import time
from xml.etree import ElementTree as ET

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from hydra.profiles import Profile
from hydra.runtime import Runtime
from hydra.session import build_labwc_config, build_viewer_argv, _viewer_environment


def wait_for(predicate, description, timeout=8):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return
        time.sleep(.05)
    raise AssertionError(description)


def counter(manager, name):
    with manager._connect(manager.paths(name)) as qmp:
        reply = qmp.command('human-monitor-command', {'command-line': 'xp /3wx 0x500'})
    values = [int(word, 16) for word in re.findall(r'0x([0-9a-fA-F]+)', reply)]
    assert len(values) == 3 and values[2] == 0x48594452, 'scratch guest did not boot'
    return values[0]


def focus_events(log):
    """Only keyboard focus, excluding pointer-enter and old window titles."""
    return re.findall(r'wl_keyboard#\d+\.(enter|leave)\(', log.read_text())


def fresh_focus_transition(source_before, target_before, source_after, target_after):
    return (len(source_after) > len(source_before) and source_after[-1] == 'leave'
            and len(target_after) > len(target_before) and target_after[-1] == 'enter')


def measure_hidden_interval(read_counter, hidden_log, visible_log, duration=.5):
    hidden_focus, visible_focus = focus_events(hidden_log), focus_events(visible_log)
    assert hidden_focus and hidden_focus[-1] == 'leave', 'sampled guest is not hidden'
    assert visible_focus and visible_focus[-1] == 'enter', 'other guest is not focused'

    def unchanged():
        assert (focus_events(hidden_log) == hidden_focus and
                focus_events(visible_log) == visible_focus), 'focus changed during hidden interval'

    start = time.monotonic()
    before = read_counter()
    unchanged()
    while time.monotonic() - start < duration:
        time.sleep(min(.05, duration))
        unchanged()
    after = read_counter()
    unchanged()
    assert after != before, 'hidden guest stopped working during continuously hidden interval'
    return {'before': before, 'after': after,
            'elapsed_seconds': time.monotonic() - start, 'focus_unchanged': True}


def classify_compositor_errors(log):
    expected = []
    for line in log.splitlines():
        if '[ERROR]' not in line:
            continue
        if re.search(r'\[xwayland/sockets\.c:\d+\] Failed to bind socket '
                     r'@?/tmp/\.X11-unix/X\d+: Address already in use$', line):
            expected.append(line)
        else:
            raise AssertionError('Unexpected compositor error: ' + line)
    return expected


def main():
    cli = argparse.ArgumentParser(description=__doc__)
    cli.add_argument('--viewer', default=os.environ.get('HYDRA_VIEWER', 'remote-viewer'))
    cli.add_argument('--labwc', default=os.environ.get('HYDRA_LABWC', 'labwc'))
    cli.add_argument('--keyboard', help='compiled dev/input-probe.c; built automatically when omitted')
    args = cli.parse_args()
    for name in ('viewer', 'labwc'):
        value = shutil.which(getattr(args, name))
        if not value:
            cli.error(f'missing {name}: {getattr(args, name)}')
        setattr(args, name, value)
    source = Path(__file__).resolve().parent
    if not args.keyboard:
        args.keyboard = subprocess.check_output([str(source / 'build-input-probe.sh')], text=True).strip()
    if not os.access('/dev/kvm', os.R_OK | os.W_OK):
        cli.error('KVM access required for this integration check')
    root = Path(tempfile.mkdtemp(prefix='hydra-smoke-', dir='/tmp'))
    print(f'Evidence: {root}', flush=True)
    xdg = root / 'xdg'
    xdg.mkdir(mode=0o700)
    config = root / 'labwc'
    config.mkdir()
    recovery = root / 'recovery'
    unreserved = root / 'unreserved'
    names = ['alpha', 'beta']
    rc = ET.fromstring(build_labwc_config(names, [sys.executable, '-c',
                           'from pathlib import Path; Path(' + repr(str(recovery)) + ').touch()'], initial='alpha'))
    # Test-only controls prove native labwc inhibition is active; production
    # config has no ToggleKeybinds or ordinary host shortcut.
    keyboard = rc.find('keyboard')
    binding = ET.SubElement(keyboard, 'keybind', key='C-A-S-i')
    ET.SubElement(binding, 'action', name='ToggleKeybinds')
    binding = ET.SubElement(keyboard, 'keybind', key='F12')
    ET.SubElement(binding, 'action', name='Execute', command=shutil.which('touch') + ' ' + str(unreserved))
    ET.ElementTree(rc).write(config / 'rc.xml')
    (config / 'autostart').write_text('#!/bin/sh\n')
    (config / 'environment').touch()
    env = dict(os.environ, XDG_RUNTIME_DIR=str(xdg), WLR_BACKENDS='headless',
               WLR_RENDERER='pixman', WLR_XWAYLAND='/bin/false', XDG_DATA_DIRS='/nonexistent',
               GSETTINGS_BACKEND='memory', LABWC_UPDATE_ACTIVATION_ENV='0')
    for key in ('DISPLAY', 'WAYLAND_DISPLAY'):
        env.pop(key, None)
    manager = Runtime(root / 'config', root / 'state', root / 'run')
    boot_object, boot = root / 'boot.o', root / 'boot.bin'
    subprocess.run(['as', '--32', str(source / 'guest-probe.S'), '-o', str(boot_object)], check=True)
    subprocess.run(['objcopy', '-j', '.text', '-O', 'binary', str(boot_object), str(boot)], check=True)
    raw = root / 'boot.raw'
    raw.write_bytes(boot.read_bytes() + bytes(8 * 1024 * 1024 - 512))
    media = root / 'installation.iso'
    media.touch()
    for name in names:
        manager.add(Profile(name, media if name == 'alpha' else None, memory=128, cpus=1, disk_size='8M', graphics='software',
                            width=1920, height=1080, refresh_millihz=144000))
        paths = manager.paths(name)
        manager.prepare(paths)
        subprocess.run([os.environ.get('HYDRA_QEMU_IMG', 'qemu-img'), 'convert', '-f', 'raw',
                        '-O', 'qcow2', str(raw), str(paths.disk)], check=True)
    processes, streams = [], []
    result = {'evidence': str(root), 'physical_checks_pending': True}
    try:
        cl = (root / 'labwc.log').open('w'); streams.append(cl)
        compositor = subprocess.Popen([args.labwc, '-d', '-C', str(config)], env=env, stdout=cl, stderr=cl)
        processes.append(compositor)
        wait_for(lambda: any(not p.name.endswith('.lock') for p in xdg.glob('wayland-*')), 'private display did not start')
        env['WAYLAND_DISPLAY'] = next(p.name for p in xdg.glob('wayland-*') if not p.name.endswith('.lock'))
        injector = subprocess.Popen([args.keyboard], env=env, stdin=subprocess.PIPE, text=True)
        processes.append(injector)

        def keys(commands):
            injector.stdin.write(commands + '\nwait 100\n'); injector.stdin.flush()
            time.sleep(.15)
            assert injector.poll() is None, 'private input helper exited'

        def chord(code):
            keys(f'down 29\ndown 56\ndown 42\ndown {code}\nup {code}\nup 42\nup 56\nup 29')

        def switch(key, source_log, target_log):
            source_before, target_before = focus_events(source_log), focus_events(target_log)
            chord(key)
            wait_for(lambda: fresh_focus_transition(source_before, target_before,
                                                   focus_events(source_log), focus_events(target_log)),
                     'reserved switch did not produce fresh keyboard leave/enter')
            return {'from_events': len(focus_events(source_log)),
                    'to_events': len(focus_events(target_log)), 'keycode': key}

        def attach(name, number):
            vl = (root / f'{name}-{number}.log').open('w'); streams.append(vl)
            previous = dict(os.environ)
            try:
                os.environ.update(env)
                viewer_env = _viewer_environment(root / f'{name}-{number}-settings')
            finally:
                os.environ.clear(); os.environ.update(previous)
            viewer_env['WAYLAND_DEBUG'] = '1'
            proc = subprocess.Popen(build_viewer_argv(manager.paths(name), args.viewer),
                                    env=viewer_env, stdout=vl, stderr=vl)
            processes.append(proc)
            log = root / f'{name}-{number}.log'
            wait_for(lambda: f'set_title("HydraOS: {name}")' in log.read_text(), f'{name} viewer did not connect')
            return proc, log

        for name in names:
            manager.start(name)
        media.unlink()  # Open QEMU media disappears; status/control/view still work.
        time.sleep(1)
        before = {name: counter(manager, name) for name in names}
        with manager._connect(manager.paths('alpha')) as qmp:
            children = qmp.command('qom-list', {'path': '/machine/peripheral-anon'})
            video = next(child['name'] for child in children if child['type'] == 'child<virtio-vga>')
            mode = {prop: qmp.command('qom-get', {'path': '/machine/peripheral-anon/' + video, 'property': prop})
                    for prop in ('xres', 'yres', 'refresh_rate')}
            assert mode == {'xres': 1920, 'yres': 1080, 'refresh_rate': 144000}, mode
            usb_before = qmp.command('human-monitor-command', {'command-line': 'info usb'})
        identities = {name: manager.paths(name).pid.read_text().strip() for name in names}
        alpha, alog = attach('alpha', 1)
        beta, blog = attach('beta', 1)
        wait_for(lambda: '.enter(' in alog.read_text(), 'alpha did not receive focus')
        # F12 must work before inhibition; then become suppressed in alpha.
        keys('down 88\nup 88')
        wait_for(unreserved.exists, 'ordinary test binding failed before inhibition')
        unreserved.unlink()
        chord(23)  # I: test-only ToggleKeybinds.
        keys('down 88\nup 88')
        assert not unreserved.exists(), 'keybinding inhibition did not suppress ordinary F12'
        switch(106, alog, blog)  # Right: beta focused, alpha continuously hidden.
        hidden_intervals = {'alpha': measure_hidden_interval(lambda: counter(manager, 'alpha'), alog, blog)}
        switch(105, blog, alog)  # Left: alpha focused, beta continuously hidden.
        hidden_intervals['beta'] = measure_hidden_interval(lambda: counter(manager, 'beta'), blog, alog)
        chord(1)  # Escape: production recovery binding still works.
        wait_for(recovery.exists, 'reserved recovery failed with inhibition')
        dialog = 'set_title("Select USB devices for redirection")'
        assert dialog not in alog.read_text(), 'USB selector opened automatically'
        # Hold modifiers through opening, then release inside the modal dialog.
        keys('down 29\ndown 56\ndown 42\ndown 22')
        wait_for(lambda: dialog in alog.read_text(), 'manual USB selector did not open', timeout=3)
        keys('up 22\nup 42\nup 56\nup 29')
        keys('down 1\nup 1')  # Close selector.
        post_usb_switches = [switch(106, alog, blog), switch(105, blog, alog)]
        assert alpha.poll() is None and beta.poll() is None
        assert identities == {name: manager.paths(name).pid.read_text().strip() for name in names}, 'switch restarted a guest'
        assert alog.read_text().count('set_title("HydraOS: alpha")') == 1, 'switch reconnected alpha'
        assert blog.read_text().count('set_title("HydraOS: beta")') == 1, 'switch reconnected beta'
        with manager._connect(manager.paths('alpha')) as qmp:
            usb_after = qmp.command('human-monitor-command', {'command-line': 'info usb'})
        assert usb_after == usb_before, 'opening selector attached an unselected USB device'
        alpha.terminate(); alpha.wait(timeout=5)
        assert manager.status('alpha')['status'] == 'running'
        alpha, attached = attach('alpha', 2)
        after = {name: counter(manager, name) for name in names}
        assert all(after[name] != before[name] for name in names), 'background guest stopped working'
        # Quit only beta; alpha's counter/lifecycle remains independent.
        manager.stop('beta', force=True)
        assert manager.status('alpha')['status'] == 'running'
        assert counter(manager, 'alpha') != after['alpha']
        compositor.terminate(); compositor.wait(timeout=5)
        assert manager.status('alpha')['status'] == 'running'
        manager.stop('alpha', force=True)
        manager.media('alpha', None)
        assert manager.profile('alpha').iso is None
        manager.start('alpha')
        time.sleep(1)
        disk_boot_counter = counter(manager, 'alpha')
        assert disk_boot_counter > 0
        expected_compositor_noise = classify_compositor_errors((root / 'labwc.log').read_text())
        result.update(hidden_guest_intervals=hidden_intervals, post_usb_switches=post_usb_switches,
                      expected_compositor_noise=expected_compositor_noise,
                      manual_usb_dialog=True, standard_us_keymap=True,
                      native_labwc_inhibition=True, reserved_switch_and_recovery=True,
                      qemu_survives_viewer_and_compositor=True, independent_lifecycle=True,
                      guest_counters_before=before, guest_counters_after=after,
                      guest_pids=identities, explicit_reattach=True, requested_qemu_mode=mode,
                      usb_devices_unchanged_without_selection=True,
                      missing_iso_does_not_block_views_or_control=True,
                      restart_without_original_media=True, disk_boot_counter=disk_boot_counter,
                      viewer_version=subprocess.check_output([args.viewer, '--version'], text=True).strip(),
                      guest_input_and_physical_modifiers='not proven by this boot counter workload')
        (root / 'result.json').write_text(json.dumps(result, indent=2) + '\n')
        print(json.dumps(result, indent=2))
    finally:
        for proc in reversed(processes):
            if proc.poll() is None:
                proc.terminate(); proc.wait(timeout=5)
        for stream in streams:
            stream.close()
        for name in names:
            if manager.status(name)['status'] != 'stopped':
                manager.stop(name, force=True)


if __name__ == '__main__':
    main()