#!/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-3d-windows-settle-20261011'))
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')
env['HOME']=str(root/'home')
(root/'home').mkdir()
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([os.environ.get('WCM_BINARY','/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 settings wrote config'
    config.write_text(original.replace('dim = 0.25', 'dim = 0.37') + '\n[audit-unrelated]\nkeep=external\n')
    time.sleep(.3)
    assert check.get_action_iface().do_action(0)
    wait(lambda:value().getboolean('cube','mouse_overview',fallback=False))
    current=value()
    result={'external_cube_dim_before':.37, 'cube_dim_after_wcm_save':float(current['cube']['dim']),
        'external_new_section_preserved':current.has_section('audit-unrelated')}
    print('RESULT',json.dumps(result),flush=True)
    (root/'wcm-save-result.json').write_text(json.dumps(result,indent=2))
    assert result['cube_dim_after_wcm_save']==.37 and result['external_new_section_preserved'], result
    assert '# Keep unrelated preferences.' in config.read_text()
    assert check.get_action_iface().do_action(0)
    wait(lambda:not value().getboolean('cube','mouse_overview',fallback=True))
    assert float(value()['cube']['dim'])==.37 and value()['audit-unrelated']['keep']=='external'
    print('PASS WCM saves two successive edits while preserving external settings and comments',flush=True)
finally:
    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)