#!/usr/bin/env python3
"""Open the Live Wallpaper, Screensaver and Lock Screen pages through WCM in a
headless stitched Wayfire session. Run under dbus-run-session.

WCM is the built binary in $WCM; metadata comes from $LW_METADATA (a directory
holding live-wallpaper.xml, screensaver.xml and wayfire-lock.xml)."""
import configparser
import json
import os
from pathlib import Path
import socket
import struct
import subprocess
import tempfile
import shutil
import time
import gi
gi.require_version('Atspi', '2.0')
from gi.repository import Atspi

assert os.geteuid() != 0
wcm = os.environ['WCM']
metadata = Path(os.environ['LW_METADATA'])
for name in ('live-wallpaper', 'screensaver', 'wayfire-lock'):
    assert (metadata / f'{name}.xml').is_file(), name
# Wayland socket paths are limited to 108 bytes, so keep the runtime dir short.
root = Path(tempfile.mkdtemp(prefix='wcm-lw-', dir='/tmp'))
env = os.environ.copy()
for key in ('WAYLAND_DISPLAY', 'DISPLAY', 'WAYFIRE_SOCKET', '_WAYFIRE_SOCKET', 'WAYFIRE_CONFIG_FILE',
            'WAYFIRE_DEFAULT_CONFIG_BACKEND', 'NO_AT_BRIDGE'):
    env.pop(key, None)
env.update(XDG_RUNTIME_DIR=str(root), XDG_CONFIG_HOME=str(root / 'config'), XDG_DATA_HOME=str(root / 'data'),
           XDG_CACHE_HOME=str(root / 'cache'), XDG_STATE_HOME=str(root / 'state'), WLR_BACKENDS='headless',
           WLR_HEADLESS_OUTPUTS='3', WLR_RENDERER='gles2', WLR_RENDER_DRM_DEVICE='/dev/dri/renderD128',
           WAYFIRE_PLUGIN_XML_PATH=str(metadata), GTK_USE_PORTAL='0', GIO_USE_VFS='local',
           GTK_MODULES='atk-bridge', XDG_DATA_DIRS='/usr/local/share:/usr/share')
os.chmod(root, 0o700)
config = root / 'wayfire.ini'
config.write_text('''# Keep unrelated preferences.
[core]
plugins = ipc stipc ipc-rules place grid
xwayland = false
stitch_enabled = true
stitch_outputs = HEADLESS-1 HEADLESS-2 HEADLESS-3
[grid]
toggle_scope = <super> KEY_ENTER
[wayfire-lock]
lock_output = DP-9
[output:HEADLESS-1]
mode = 1400x1000@60
position = 0,0
[output:HEADLESS-2]
mode = 1400x1000@60
position = 1400,0
[output:HEADLESS-3]
mode = 1400x1000@60
position = 2800,0
''')
original = config.read_text()
shell = root / 'shell.ini'
shell.write_text('')
compositor = subprocess.Popen(['wayfire', '-c', str(config)], env=env,
                              stdout=(root / 'wayfire.log').open('w'), stderr=subprocess.STDOUT)
ui = None


def wait(fn, timeout=20):
    until = time.monotonic() + timeout
    while time.monotonic() < until:
        assert compositor.poll() is None, f'Compositor exited; {root}'
        if ui is not None:
            assert ui.poll() is None, f'WCM exited; {root}'
        value = fn()
        if value:
            return value
        time.sleep(.05)
    raise TimeoutError(str(root))


def ipc(method, data=None):
    with socket.socket(socket.AF_UNIX) as stream:
        stream.settimeout(5)
        stream.connect(env['WAYFIRE_SOCKET'])
        payload = json.dumps({'method': method, 'data': data or {}}).encode()
        stream.sendall(struct.pack('<I', len(payload)) + payload)

        def read(n):
            value = b''
            while len(value) < n:
                block = stream.recv(n - len(value))
                assert block
                value += block
            return value
        result = json.loads(read(struct.unpack('<I', read(4))[0]))
        assert not isinstance(result, dict) or not result.get('error'), result
        return result


def nodes(start=None):
    stack = [start or Atspi.get_desktop(0)]
    while stack:
        node = stack.pop()
        if node is None:
            continue
        yield node
        stack.extend(node.get_child_at_index(i) for i in range(node.get_child_count()))


def find(name, role):
    return next((n for n in nodes() if n.get_name() == name and n.get_role() == role), None)


def showing(node):
    return node is not None and node.get_state_set().contains(Atspi.StateType.SHOWING)


def find_showing(name, role):
    return next((n for n in nodes() if n.get_name() == name and n.get_role() == role and showing(n)), None)


def control(label_name, role):
    label = find_showing(label_name, Atspi.Role.LABEL)
    if label is None:
        return None
    parent = label.get_parent()
    for _ in range(4):
        found = [n for n in nodes(parent) if n.get_role() == role]
        if len(found) == 1:
            return found[0]
        parent = parent.get_parent()
    return None


def value():
    parsed = configparser.ConfigParser(interpolation=None, strict=False)
    parsed.read(config)
    return parsed


def start_wcm(*args):
    global ui
    ui = subprocess.Popen([wcm, '--config', str(config), '--shell-config', str(shell), *args],
                          env=env, stdout=(root / 'wcm.log').open('a'), stderr=subprocess.STDOUT)


def stop_wcm():
    global ui
    if ui is not None:
        ui.terminate()
        try:
            ui.wait(timeout=5)
        except subprocess.TimeoutExpired:
            ui.kill()
            ui.wait(timeout=5)
    ui = None
    wait(lambda: not any(n.get_name() == 'wcm' for n in nodes(Atspi.get_desktop(0)) if n.get_role() == Atspi.Role.APPLICATION))


def tabs():
    return {n.get_name() for n in nodes() if n.get_role() == Atspi.Role.PAGE_TAB and showing(n.get_parent())}


passed = []
try:
    env['WAYFIRE_SOCKET'] = str(wait(lambda: next(root.glob('wayfire-*.socket'), None)))
    env['WAYLAND_DISPLAY'] = ipc('stipc/get_display')['wayland']

    # 1. Settings' "Lock screen" link lands on the Lock Screen page, with no TOML tab.
    start_wcm('--plugin=wayfire-lock')
    wait(lambda: find_showing('Lock-screen background', Atspi.Role.LABEL))
    assert 'Integration' not in tabs(), tabs()
    assert {'Background', 'Controls'} <= tabs(), tabs()
    passed.append('--plugin=wayfire-lock opens Lock Screen')

    # 2. The lock monitor combo lists stitched members and keeps the offline choice.
    tab = wait(lambda: find('Controls', Atspi.Role.PAGE_TAB))
    assert tab.get_parent().get_selection_iface().select_child(tab.get_index_in_parent())
    combo = wait(lambda: control('Lock controls monitor', Atspi.Role.COMBO_BOX))
    wait(lambda: showing(combo))
    assert combo.get_name() == 'DP-9 (disconnected)', combo.get_name()
    assert config.read_text() == original, 'Opening the page changed preferences'
    assert combo.get_action_iface().do_action(0)
    names = wait(lambda: (lambda s: s if {'HEADLESS-1', 'HEADLESS-2', 'HEADLESS-3'} <= s else None)(
        {n.get_name() for n in nodes() if n.get_role() == Atspi.Role.MENU_ITEM and showing(n)}))
    assert 'VIRTUAL-1' not in names, names
    choice = find_showing('HEADLESS-2', Atspi.Role.MENU_ITEM)
    assert choice.get_action_iface().do_action(0)
    wait(lambda: value().get('wayfire-lock', 'lock_output', fallback='') == 'HEADLESS-2')
    wait(lambda: ipc('wayfire/get-config-option', {'option': 'wayfire-lock/lock_output'})['value'] == 'HEADLESS-2')
    assert value()['grid']['toggle_scope'] == '<super> KEY_ENTER'
    passed.append('lock monitor lists HEADLESS-1..3, saves wayfire-lock/lock_output')
    stop_wcm()

    # 3. Main page: three Desktop entries, no Digital Rain button; each page opens.
    for button, label in (('Live Wallpaper', 'Enable live wallpaper'),
                          ('Screensaver', 'Focused monitor idle delay'),
                          ('Lock Screen', 'Lock-screen background')):
        start_wcm()
        page = wait(lambda: find(button, Atspi.Role.PUSH_BUTTON))
        assert find('Digital Rain', Atspi.Role.PUSH_BUTTON) is None
        heading = page.get_parent()
        while heading is not None and not any(n.get_name() == 'Desktop' and n.get_role() == Atspi.Role.LABEL
                                              for n in nodes(heading)):
            heading = heading.get_parent()
        assert heading is not None and heading.get_role() != Atspi.Role.FRAME, f'{button} not under Desktop'
        page.get_component_iface().scroll_to(Atspi.ScrollType.ANYWHERE)
        wait(lambda: showing(page))
        assert page.get_action_iface().do_action(0)
        wait(lambda: find_showing(label, Atspi.Role.LABEL))
        assert 'Integration' not in tabs(), (button, tabs())
        passed.append(f'{button} page opens from Desktop')
        stop_wcm()

    # 4. The retired plugin names are no longer aliases: they open the main page.
    for retired in ('digitalrain', 'ambient-background'):
        start_wcm('--plugin=' + retired)
        wait(lambda: find('Live Wallpaper', Atspi.Role.PUSH_BUTTON))
        time.sleep(1)
        assert find_showing('Lock-screen background', Atspi.Role.LABEL) is None, retired
        stop_wcm()
    passed.append('--plugin=digitalrain / ambient-background are not redirected')
    for line in passed:
        print('PASS', line, flush=True)
    print(f'{len(passed)}/{len(passed)} passed', flush=True)
finally:
    if ui is not None:
        ui.terminate()
        try:
            ui.wait(timeout=5)
        except subprocess.TimeoutExpired:
            ui.kill()
            ui.wait(timeout=5)
    compositor.terminate()
    try:
        compositor.wait(timeout=5)
    except subprocess.TimeoutExpired:
        compositor.kill()
        compositor.wait(timeout=5)
    if len(passed) == 6:
        shutil.rmtree(root, ignore_errors=True)
    else:
        print('artifacts:', root)