#!/usr/bin/env python3
"""Live input bridge — the interactive half of the client-GPU architecture.

Browser sends control events over WebSocket; we inject them into the live
Warcraft III window via XTEST (the SAME mechanism production streaming uses,
streaming/server.py). This proves input->game reaction is preserved verbatim
when the pixel path moves to client-GPU rendering.

The frame RETURN path (game d3d9 commands -> browser) is the recorder-proxy
DLL, the one remaining production component; this bridge streams periodic
Xvfb screenshots only as visual proof the game reacted.
"""
import asyncio
import io
import json
import subprocess
import sys

import websockets
from PIL import Image
from Xlib import X, display
from Xlib.ext import xtest
from Xlib import XK

DISPLAY = ":9"
KEYMAP = {"ArrowUp": "Up", "ArrowDown": "Down", "ArrowLeft": "Left",
          "ArrowRight": "Right", "p": "p", "=": "equal", "-": "minus",
          " ": "space", "Escape": "Escape"}


def game_window(conn):
    root = conn.screen().root
    for outer in root.query_tree().children:
        try:
            if outer.get_wm_name() != "Warcraft III":
                continue
        except Exception:
            continue
        og = outer.get_geometry()
        for child in outer.query_tree().children:
            g = child.get_geometry()
            if g.width == og.width and g.height == og.height:
                return child, og
        return outer, og
    return root, root.get_geometry()


class Injector:
    def __init__(self):
        self.disp = display.Display(DISPLAY)
        self.win, self.geo = game_window(self.disp)
        print(f"game window {self.geo.width}x{self.geo.height}", flush=True)

    def move(self, nx, ny):
        x = int(nx * self.geo.width)
        y = int(ny * self.geo.height)
        xtest.fake_input(self.disp, X.MotionNotify, 0, X.CurrentTime,
                         self.disp.screen().root, x, y)
        self.disp.sync()

    def click(self, nx, ny, button=1):
        self.move(nx, ny)
        xtest.fake_input(self.disp, X.ButtonPress, button)
        xtest.fake_input(self.disp, X.ButtonRelease, button)
        self.disp.sync()

    def key(self, name, down=None):
        keysym = XK.string_to_keysym(KEYMAP.get(name, name))
        if keysym == 0:
            return
        kc = self.disp.keysym_to_keycode(keysym)
        if down is None:
            xtest.fake_input(self.disp, X.KeyPress, kc)
            xtest.fake_input(self.disp, X.KeyRelease, kc)
        else:
            xtest.fake_input(self.disp,
                             X.KeyPress if down else X.KeyRelease, kc)
        self.disp.sync()

    def screenshot(self):
        g = self.win.get_geometry()
        raw = self.win.get_image(0, 0, g.width, g.height, X.ZPixmap,
                                 0xFFFFFFFF)
        img = Image.frombytes("RGB", (g.width, g.height), raw.data,
                              "raw", "BGRX")
        buf = io.BytesIO()
        img.save(buf, "JPEG", quality=70)
        return buf.getvalue()


async def handler(ws):
    inj = Injector()
    print("client connected", flush=True)

    async def push_frames():
        while True:
            try:
                await ws.send(inj.screenshot())
            except Exception:
                return
            await asyncio.sleep(0.2)          # 5 fps proof stream

    pusher = asyncio.create_task(push_frames())
    try:
        async for msg in ws:
            if isinstance(msg, bytes):
                continue
            e = json.loads(msg)
            t = e.get("t")
            if t == "m":
                inj.move(e["x"], e["y"])
            elif t == "c":
                inj.click(e["x"], e["y"], e.get("b", 1))
            elif t == "k":
                inj.key(e["k"], e.get("d"))
            elif t == "tap":
                inj.key(e["k"])
            print("inject", e, flush=True)
    finally:
        pusher.cancel()


async def main():
    async with websockets.serve(handler, "0.0.0.0", 8145, max_size=None):
        print("bridge on :8145", flush=True)
        await asyncio.Future()

asyncio.run(main())