"""Private Xvfb screen capture and XTest input bridge for the native session."""
import ctypes as C
import json
import os
from pathlib import Path
import sys
import time
from PIL import ImageGrab
from PIL import Image

root=Path(sys.argv[1]);inbox=root/'events';inbox.mkdir(exist_ok=True)
request=json.loads((root/'request.json').read_text())
width=request.get('desktop',{}).get('width',1280);height=request.get('desktop',{}).get('height',900)
frame_interval=1/20 if request.get('desktop',{}).get('preserveLayout') else 0.4
pen_cursor=None
x=C.CDLL('libX11.so.6');xt=C.CDLL('libXtst.so.6')
x.XOpenDisplay.argtypes=[C.c_char_p];x.XOpenDisplay.restype=C.c_void_p
x.XStringToKeysym.argtypes=[C.c_char_p];x.XStringToKeysym.restype=C.c_ulong
x.XKeysymToKeycode.argtypes=[C.c_void_p,C.c_ulong];x.XKeysymToKeycode.restype=C.c_uint
x.XkbKeycodeToKeysym.argtypes=[C.c_void_p,C.c_ubyte,C.c_uint,C.c_uint];x.XkbKeycodeToKeysym.restype=C.c_ulong
x.XFlush.argtypes=[C.c_void_p]
xt.XTestFakeMotionEvent.argtypes=[C.c_void_p,C.c_int,C.c_int,C.c_int,C.c_ulong]
xt.XTestFakeButtonEvent.argtypes=[C.c_void_p,C.c_uint,C.c_int,C.c_ulong]
xt.XTestFakeKeyEvent.argtypes=[C.c_void_p,C.c_uint,C.c_int,C.c_ulong]
display=x.XOpenDisplay(None)
if not display:raise RuntimeError('NATIVE_DISPLAY_UNAVAILABLE')
class ModifierMap(C.Structure):
    _fields_=[('max_keypermod',C.c_int),('modifiermap',C.POINTER(C.c_ubyte))]
x.XGetModifierMapping.argtypes=[C.c_void_p];x.XGetModifierMapping.restype=C.POINTER(ModifierMap)
x.XFreeModifiermap.argtypes=[C.POINTER(ModifierMap)]
x.XkbLockModifiers.argtypes=[C.c_void_p,C.c_uint,C.c_uint,C.c_uint]
mapping=x.XGetModifierMapping(display);num_mask=0
if mapping:
    num_code=x.XKeysymToKeycode(display,x.XStringToKeysym(b'Num_Lock'));count=mapping.contents.max_keypermod
    for index in range(8):
        if any(mapping.contents.modifiermap[index*count+i]==num_code for i in range(count)):num_mask|=1<<index
    x.XFreeModifiermap(mapping)
class CursorImage(C.Structure):
    _fields_=[('x',C.c_short),('y',C.c_short),('width',C.c_ushort),('height',C.c_ushort),
              ('xhot',C.c_ushort),('yhot',C.c_ushort),('serial',C.c_ulong),('pixels',C.POINTER(C.c_ulong)),
              ('atom',C.c_ulong),('name',C.c_char_p)]
try:
    fixes=C.CDLL('libXfixes.so.3');fixes.XFixesGetCursorImage.argtypes=[C.c_void_p]
    fixes.XFixesGetCursorImage.restype=C.POINTER(CursorImage);x.XFree.argtypes=[C.c_void_p]
except OSError:fixes=None
(root/'screen-meta.json').write_text(json.dumps({'cursorEmbedded':fixes is not None}))
def draw_cursor(picture):
    if fixes is None:return
    pointer=fixes.XFixesGetCursorImage(display)
    if not pointer:return
    try:
        value=pointer.contents
        if not value.width or not value.height or value.width*value.height>65536:return
        rgba=bytearray()
        for index in range(value.width*value.height):
            pixel=value.pixels[index];rgba.extend(((pixel>>16)&255,(pixel>>8)&255,pixel&255,(pixel>>24)&255))
        cursor=Image.frombytes('RGBA',(value.width,value.height),bytes(rgba))
        px,py=pen_cursor if pen_cursor is not None else (value.x,value.y)
        picture.paste(cursor,(px-value.xhot,py-value.yhot),cursor)
    finally:x.XFree(pointer)
keys={'Enter':'Return','Escape':'Escape','Backspace':'BackSpace','Tab':'Tab','Space':'space',
      'ArrowUp':'Up','ArrowDown':'Down','ArrowLeft':'Left','ArrowRight':'Right','Delete':'Delete',
      'Home':'Home','End':'End','PageUp':'Prior','PageDown':'Next','Minus':'minus','Equal':'equal',
      'Period':'period','Comma':'comma','Slash':'slash','Semicolon':'semicolon','Quote':'apostrophe',
      'BracketLeft':'bracketleft','BracketRight':'bracketright','Backslash':'backslash',
      'ShiftLeft':'Shift_L','ShiftRight':'Shift_R','ControlLeft':'Control_L','ControlRight':'Control_R',
      'AltLeft':'Alt_L','AltRight':'Alt_R','MetaLeft':'Super_L','MetaRight':'Super_R'}
keys.update({'NumpadEnter':'KP_Enter','NumpadDecimal':'KP_Decimal','NumpadDivide':'KP_Divide',
             'NumpadMultiply':'KP_Multiply','NumpadSubtract':'KP_Subtract','NumpadAdd':'KP_Add',
             'Backquote':'grave','Insert':'Insert','CapsLock':'Caps_Lock','NumLock':'Num_Lock'})
for letter in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ':keys['Key'+letter]=letter.lower()
for digit in '0123456789':keys['Digit'+digit]=digit;keys['Numpad'+digit]='KP_'+digit
for number in range(1,13):keys['F'+str(number)]='F'+str(number)
pressed={};buttons=set();last_frame=0;last_pixels=None
def key(code,down,logical=None,caps=False):
    symbol=keys.get(code)
    if code not in pressed and not symbol and not (logical and len(logical)==1 and ord(logical)<256):return
    if code in pressed:number=pressed[code]
    else:
        keysym=ord(logical) if logical and len(logical)==1 and ord(logical)<256 and not code.startswith('Numpad') else x.XStringToKeysym(symbol.encode())
        number=x.XKeysymToKeycode(display,keysym)
    shift=0
    if down and number and logical and len(logical)==1 and ord(logical)<128 and not code.startswith('Numpad') and not any(code in pressed for code in ('ShiftLeft','ShiftRight')):
        needed=logical.isupper()!=caps if logical.isalpha() else x.XkbKeycodeToKeysym(display,number,0,1)==ord(logical) and x.XkbKeycodeToKeysym(display,number,0,0)!=ord(logical)
        if needed:shift=x.XKeysymToKeycode(display,x.XStringToKeysym(b'Shift_L'));xt.XTestFakeKeyEvent(display,shift,1,1)
    # GHOST_SystemX11 rejects a same-key press at the preceding release timestamp.
    # XTEST uses server milliseconds; distinct presses must advance that clock tick.
    if number:xt.XTestFakeKeyEvent(display,number,int(down),1 if down and code not in pressed else 0)
    if shift:xt.XTestFakeKeyEvent(display,shift,0,0)
    if down:pressed[code]=number
    else:pressed.pop(code,None)
def event(item):
    global pen_cursor,pen_fallback
    kind=item['type']
    if kind in ('move','button'):
        pen_cursor=None
        xt.XTestFakeMotionEvent(display,-1,max(0,min(width-1,item['x'])),max(0,min(height-1,item['y'])),0)
    if kind=='button':
        button=item['button'];xt.XTestFakeButtonEvent(display,button,int(item['down']),0)
        if item['down']:buttons.add(button)
        else:buttons.discard(button)
    elif kind=='key':
        locks=item.get('locks')
        if locks and item['code'] not in ('CapsLock','NumLock'):
            x.XkbLockModifiers(display,0x0100,2|num_mask,(2 if locks['caps'] else 0)|(num_mask if locks['num'] else 0))
        key(item['code'],item['down'],item.get('key'),bool((locks or {}).get('caps')))
    elif kind=='wheel':
        button=(6 if item['delta']<0 else 7) if item.get('horizontal') else (4 if item['delta']<0 else 5)
        for _ in range(min(5,abs(item['delta']))):
            xt.XTestFakeButtonEvent(display,button,1,0);xt.XTestFakeButtonEvent(display,button,0,0)
    elif kind=='release':
        for code in list(pressed):key(code,False)
        for button in list(buttons):xt.XTestFakeButtonEvent(display,button,0,0)
        buttons.clear()
        pen_fallback=False
    x.XFlush(display)
def publish(file,value):
    temporary=file.with_suffix('.tmp');temporary.write_text(json.dumps(value));temporary.replace(file)
current=None;pen_pending=None;pen_sequence=0;pen_active=None;pen_fallback=False
def pen_mouse(item):
    event({'type':'move','x':item['x'],'y':item['y']})
    if item['action']!='move':event({'type':'button','x':item['x'],'y':item['y'],'button':item['button'],'down':item['action']=='press'})
while not (root/'stop').exists():
    if pen_pending:
        reply=root/'pen-output.json'
        if reply.exists():
            result=json.loads(reply.read_text())
            if result.get('id')==pen_pending['id']:
                if result.get('error'):publish(root/'input-error.json',{'error':result['error']});pen_active=None
                else:
                    for sample,handled in zip(pen_pending['samples'],result['handled']):
                        if handled:
                            pen_cursor=(sample['x'],sample['y'])
                            pen_active=sample if sample['action']!='release' and (pen_active or sample['action']=='press') else None
                        else:
                            pen_mouse(sample)
                            if sample['action']=='press':pen_fallback=True
                            elif sample['action']=='release':pen_fallback=False;pen_active=None
                pen_pending=None;reply.unlink()
        if pen_pending and time.monotonic()-pen_pending['started']>10:
            publish(root/'input-error.json',{'error':'NATIVE_TABLET_INPUT_TIMEOUT'})
            raise RuntimeError('NATIVE_TABLET_INPUT_TIMEOUT')
    if current is None and pen_pending is None:
        file=next(iter(sorted(inbox.glob('*.json'))),None)
        if file:current={'file':file,'items':json.loads(file.read_text()),'index':0}
    if current and pen_pending is None:
        while current['index']<len(current['items']):
            item=current['items'][current['index']]
            if item['type']=='release' and pen_active:
                samples=[dict(pen_active,action='release',pressure=0.0)]
            elif item['type']=='pen' and not pen_fallback:
                samples=[]
                while current['index']<len(current['items']) and current['items'][current['index']]['type']=='pen':
                    sample=current['items'][current['index']]
                    if samples and (sample['action']!='move' or samples[0]['action']!='move'):break
                    samples.append(sample);current['index']+=1
                if samples[0]['action']=='press':event({'type':'move','x':samples[0]['x'],'y':samples[0]['y']})
            else:
                if item['type']=='pen':
                    pen_mouse(item)
                    if item['action']=='release':pen_fallback=False
                else:event(item)
                current['index']+=1;continue
            pen_sequence+=1;pen_pending={'id':pen_sequence,'samples':samples,'started':time.monotonic()}
            publish(root/'pen-input.json',{'id':pen_sequence,'samples':samples});break
        if pen_pending is None and current['index']==len(current['items']):
            file=current['file']
            if file.stem.isdigit():publish(root/'input-ack.json',{'sequence':int(file.stem),'at':time.time()})
            file.unlink();current=None
    if time.monotonic()-last_frame>=frame_interval:
        picture=ImageGrab.grab(xdisplay=os.environ['DISPLAY'])
        draw_cursor(picture)
        pixels=picture.tobytes()
        if pixels!=last_pixels:
            picture.save(root/'screen.tmp',format='PNG',compress_level=1);(root/'screen.tmp').replace(root/'screen.png')
            last_pixels=pixels
        last_frame=time.monotonic()
    time.sleep(0.015)