#!/usr/bin/env python3
import json, os, socket, struct, subprocess, tempfile, time
from pathlib import Path
stage=Path(os.environ.get('CUBE_TEST_STAGE', '/var/tmp/wayfire-cube-repeat'))
runtime=Path(tempfile.mkdtemp(prefix='cube-repeat-'))
env=os.environ.copy()
for k in ('WAYLAND_DISPLAY','DISPLAY','WAYFIRE_SOCKET','_WAYFIRE_SOCKET','WAYFIRE_CONFIG_FILE','WAYFIRE_DEFAULT_CONFIG_BACKEND'):
    env.pop(k,None)
env.update(XDG_RUNTIME_DIR=str(runtime),XDG_DATA_HOME=str(runtime/'data'),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'))
config=runtime/'wayfire.ini'
config.write_text('''[core]
plugins = ipc ipc-rules stipc cube
xwayland = false
stitch_enabled = true
stitch_outputs = HEADLESS-1 HEADLESS-2 HEADLESS-3
vwidth = 16
vheight = 1
[cube]
rotate_left = <super> KEY_LEFT | <alt> <ctrl> KEY_LEFT
rotate_right = <super> KEY_RIGHT | <alt> <ctrl> KEY_RIGHT
[autostart]
autostart_wf_shell = false
[output:HEADLESS-1]
mode = 320x240@60
position = 0,0
[output:HEADLESS-2]
mode = 320x240@60
position = 320,0
[output:HEADLESS-3]
mode = 320x240@60
position = 640,0
''')
log=(runtime/'wayfire.log').open('w')
process=subprocess.Popen(['/usr/bin/wayfire','-c',str(config)],env=env,stdout=log,stderr=log)
def ipc(method,data=None):
    with socket.socket(socket.AF_UNIX) as s:
        s.settimeout(5);s.connect(str(sock))
        b=json.dumps({'method':method,'data':data or {}}).encode();s.sendall(struct.pack('<I',len(b))+b)
        def read(n):
            b=b''
            while len(b)<n:
                c=s.recv(n-len(b))
                if not c:raise RuntimeError('IPC disconnected')
                b+=c
            return b
        r=json.loads(read(struct.unpack('<I',read(4))[0]))
        assert not isinstance(r,dict) or not r.get('error'),r
        return r
def key(k,state):ipc('stipc/feed_key',{'key':k,'state':state})
def options(**kw):ipc('wayfire/set-config-options',kw)
def ws():return ipc('window-rules/list-outputs')[0]['workspace']['x']
def watch(seconds):
    last=ws(); changes=[];start=time.monotonic()
    while time.monotonic()-start<seconds:
        n=ws()
        if n!=last:changes.append((time.monotonic()-start,n));last=n
        time.sleep(.008)
    return changes
def press(k='KEY_RIGHT'):key('KEY_LEFTMETA',True);key(k,True)
def release(k='KEY_RIGHT'):key(k,False);key('KEY_LEFTMETA',False)
def stable(seconds=.8):assert not watch(seconds),'rotation continued after release'
try:
    print('Runtime:',runtime,flush=True)
    deadline=time.monotonic()+10
    while not (sockets:=list(runtime.glob('wayfire-*.socket'))):
        assert process.poll() is None,f'Compositor exited: {runtime}/wayfire.log'
        assert time.monotonic()<deadline
        time.sleep(.05)
    sock=sockets[0];time.sleep(.5)
    print('Outputs:',ipc('window-rules/list-outputs'),flush=True)
    # Start with the old cube resident, then load the candidate at a fresh path.
    plugins=runtime/'data/wayfire/plugins';plugins.mkdir(parents=True)
    (plugins/'libcube.so').symlink_to(stage/'cube-next.so')
    (stage/'plugins/libcube.so').unlink()
    ipc('wayfire/reload-plugins');time.sleep(.3)
    assert ipc('wayfire/get-config-option',{'option':'cube/repeat_enabled'})['value']
    print('PASS plugin hot-load over resident original cube',flush=True)
    before=ws();press();assert ws()==(before+1)%16
    assert not watch(.40),'repeated before initial pause'
    release();stable()
    print('PASS immediate first step and single tap',flush=True)
    press();changes=watch(4.7);release();stable()
    gaps=[changes[i+1][0]-changes[i][0] for i in range(len(changes)-1)]
    print('Acceleration changes:',changes,'intervals:',gaps,flush=True)
    assert len(gaps)>8 and .50<changes[0][0]<.80
    assert .48<gaps[0]<.75 and .10<gaps[-1]<.24 and gaps[0]>gaps[-1]*2
    print('PASS slow-to-fast acceleration and arrow release',flush=True)
    press();watch(1.4);key('KEY_LEFTMETA',False);stable();key('KEY_RIGHT',False)
    print('PASS modifier release stops repeat',flush=True)
    press();assert not watch(.40);release();stable()
    print('PASS next hold starts with a fresh delay',flush=True)
    before=ws()
    for k in ['KEY_LEFTCTRL','KEY_LEFTALT','KEY_RIGHT']:key(k,True)
    assert ws()==(before+1)%16
    assert watch(.8), 'alternate binding did not repeat'
    for k in ['KEY_RIGHT','KEY_LEFTALT','KEY_LEFTCTRL']:key(k,False)
    stable()
    print('PASS Ctrl+Alt alternate binding also repeats',flush=True)
    press();watch(.8);key('KEY_A',True);stable();key('KEY_A',False);release()
    print('PASS unrelated key stops the held shortcut',flush=True)
    options(**{'cube/repeat_acceleration':False,'cube/repeat_interval':300,'cube/repeat_delay':250})
    press('KEY_LEFT');changes=watch(1.6);release('KEY_LEFT');stable()
    gaps=[changes[i+1][0]-changes[i][0] for i in range(len(changes)-1)]
    assert len(gaps)>=3 and all(.24<g<.38 for g in gaps),gaps
    assert all(changes[i+1][1]==(changes[i][1]-1)%16 for i in range(len(changes)-1))
    print('PASS constant repeat option and reverse direction',flush=True)
    options(**{'cube/repeat_enabled':False});before=ws();press();assert ws()==(before+1)%16
    stable(1);release()
    print('PASS repeat disabled',flush=True)
    options(**{'cube/repeat_enabled':True});before=ws();ipc('cube/rotate_right');assert ws()==(before+1)%16;stable(1)
    print('PASS IPC remains a single step',flush=True)
    press();watch(.8)
    config.write_text(config.read_text().replace(' stipc cube',' stipc'))
    time.sleep(.3);stable();release()
    config.write_text(config.read_text().replace(' stipc\n',' stipc cube\n'))
    time.sleep(.4);press();watch(.8);release();stable()
    print('PASS unloading while held cancels timer and reload works',flush=True)
    print('ALL CUBE REPEAT CHECKS PASSED',flush=True)
finally:
    process.terminate();process.wait(timeout=10);log.close()
    if not (stage/'plugins/libcube.so').exists():
        import shutil
        shutil.copyfile(stage/'cube-original.so',stage/'plugins/libcube.so')