#!/usr/bin/env python3
"""Toggle Cube's overview through WCM on a private session and config."""
import configparser
import json
import os
from pathlib import Path
import socket
import struct
import subprocess
import tempfile
import time
import gi
gi.require_version('Atspi', '2.0')
from gi.repository import Atspi

assert os.geteuid()!=0
stage=Path(os.environ.get('CUBE_OVERVIEW_STAGE','/var/tmp/wayfire-cube-overview-20261009'))
root=Path(tempfile.mkdtemp(prefix='wcm-cube-overview-'))
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_PATH=str(stage/'plugins'),WAYFIRE_PLUGIN_XML_PATH=str(stage/'metadata'),
           GTK_USE_PORTAL='0',GIO_USE_VFS='local',GTK_MODULES='atk-bridge',
           XDG_DATA_DIRS='/usr/local/share:/usr/share')
config=root/'wayfire.ini'
config.write_text(f'''# Keep unrelated preferences.
[core]
plugins = ipc stipc ipc-rules {stage}/plugins/libcube.so place grid
xwayland = false
vwidth = 4
vheight = 3
stitch_enabled = true
stitch_outputs = HEADLESS-1 HEADLESS-2 HEADLESS-3
[cube]
activate = <ctrl> <alt> BTN_LEFT
dim = 0.25
[grid]
toggle_scope = <super> KEY_ENTER
[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([str(stage/'wayfire'),'-c',str(config)],env=env,
    stdout=(root/'wayfire.log').open('w'),stderr=subprocess.STDOUT)
ui=None

def wait(fn,timeout=15):
    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'])
        assert struct.unpack('3i',stream.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))[:2]==(compositor.pid,os.geteuid())
        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():
    stack=[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 overview_check(name='Show all workspace rows'):
    # Generic WCM options place a separate label beside an unnamed checkbox.
    label=find(name,Atspi.Role.LABEL)
    if label is None:return None
    container=label.get_parent()
    for _ in range(4):
        stack=[container];checks=[]
        while stack:
            node=stack.pop()
            if node.get_role()==Atspi.Role.CHECK_BOX:checks.append(node)
            stack.extend(node.get_child_at_index(i) for i in range(node.get_child_count()))
        if len(checks)==1:return checks[0]
        container=container.get_parent()
    return None

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

try:
    env['WAYFIRE_SOCKET']=str(wait(lambda:next(root.glob('wayfire-*.socket'),None)))
    env['WAYLAND_DISPLAY']=ipc('stipc/get_display')['wayland']
    ui=subprocess.Popen(['/usr/bin/wcm','--config',str(config),'--shell-config',str(shell),'--plugin=cube'],
        env=env,stdout=(root/'wcm.log').open('w'),stderr=subprocess.STDOUT)
    check=wait(overview_check)
    tab=find('Mouse overview',Atspi.Role.PAGE_TAB)
    if tab is not None:
        assert tab.get_parent().get_selection_iface().select_child(tab.get_index_in_parent())
    wait(lambda:showing(check))
    time.sleep(.3)
    assert config.read_text()==original,'Opening Cube settings wrote the config'
    assert not check.get_state_set().contains(Atspi.StateType.CHECKED)
    subprocess.run(['grim','-o','VIRTUAL-1',str(root/'cube-option.png')],env=env,check=True)
    print('PASS Cube checkbox appears with the default disabled and no config writes',flush=True)
    for enabled in (True,False,True):
        assert check.get_action_iface().do_action(0)
        wait(lambda:value().getboolean('cube','mouse_overview',fallback=not enabled)==enabled)
        wait(lambda:ipc('wayfire/get-config-option',{'option':'cube/mouse_overview'})['value']==str(enabled).lower())
        assert value()['core']['plugins'].split()==['ipc','stipc','ipc-rules',str(stage/'plugins/libcube.so'),'place','grid']
        assert set(value()['cube']['activate'].split())=={'<ctrl>','<alt>','BTN_LEFT'}
        assert float(value()['cube']['dim'])==.25
        assert value()['grid']['toggle_scope']=='<super> KEY_ENTER'
    print('PASS WCM saves and applies enable/disable changes while preserving plugin path and other preferences',flush=True)
    rest=wait(lambda:overview_check('Render other faces at rest'))
    tab=find('General',Atspi.Role.PAGE_TAB)
    if tab is not None:
        assert tab.get_parent().get_selection_iface().select_child(tab.get_index_in_parent())
    wait(lambda:showing(rest))
    assert not rest.get_state_set().contains(Atspi.StateType.CHECKED)
    for enabled in (True,False,True):
        assert rest.get_action_iface().do_action(0)
        wait(lambda:value().getboolean('cube','render_at_rest',fallback=not enabled)==enabled)
        wait(lambda:ipc('wayfire/get-config-option',{'option':'cube/render_at_rest'})['value']==str(enabled).lower())
        assert value()['cube'].getboolean('mouse_overview')
        assert value()['core']['plugins'].split()==['ipc','stipc','ipc-rules',str(stage/'plugins/libcube.so'),'place','grid']
    print('PASS WCM applies resting-face rendering without changing mouse overview or plugin paths',flush=True)
    print('ALL CUBE WCM CHECKS PASSED',flush=True)
finally:
    try:(root/'accessibility.txt').write_text('\n'.join(str(n.get_role())+' '+n.get_name() for n in nodes()))
    except Exception:pass
    if ui is not None and ui.poll() is None:ui.terminate();ui.wait(timeout=5)
    compositor.terminate();compositor.wait(timeout=10)
    print('Artifacts:',root,flush=True)