#!/usr/bin/env python3
# rikkiti-msg — send a command, or run a query, against the Rikkiti compositor IPC
# (the socket at $RIKKITISOCK; i3-ipc framing). Like swaymsg.
#
#   rikkiti-msg output HDMI-A-1 scale 1.5     # run a compositor command
#   rikkiti-msg workspace 3
#   rikkiti-msg seat click left 640 360       # synthesise input (see the comp)
#   rikkiti-msg get-outputs                   # query → prints the JSON reply
#   rikkiti-msg get-workspaces
#   rikkiti-msg get-windows
#
# Agent conveniences (built on get-windows; no protocol change):
#   rikkiti-msg get-window <app_id>|--focused          # the matching window object(s)
#   rikkiti-msg geom <app_id>|--focused [--frame]      # a `grim -g` geometry string
#                                                      # --frame includes the SSD titlebar
#   rikkiti-msg wait-window <app_id>|--focused [secs]  # block until it appears → rect
#   rikkiti-msg capture window <app_id>|--focused [--max N] [-o PATH]
#                                                      # comp-native window capture → PNG.
#                                                      # Reliable + NON-INTRUSIVE: grabs the
#                                                      # window even when occluded (no raise/focus).
#   rikkiti-msg capture output <name>|--focused [--max N] [-o PATH]  # one monitor (full scene)
#   rikkiti-msg capture region <x,y wxh>            [--max N] [-o PATH]  # a rectangle (may span outputs)
#   rikkiti-msg capture all                         [--max N] [-o PATH]  # the whole desktop
#   rikkiti-msg seat click-in <app_id> [button] <rx> <ry>   # click at window-relative px
#   rikkiti-msg seat move-in/dblclick-in/drag-in <app_id> … # (same, other ops)
import socket, struct, sys, os, json, time, glob, mmap, zlib

QUERIES = {"get-workspaces": 1, "get-windows": 2, "get-outputs": 3}
CAPTURE = 6  # RIK_MSG_CAPTURE — full-res window capture, reply = JSON + 1 memfd (docs/38)

def _candidate_socks():
    """Comp IPC sockets to try, newest first. RIKKITISOCK wins when exported; else
    every rikkiti-ipc-*.sock in the runtime dir. The comp normally exports RIKKITISOCK
    to its children, but tools launched outside that env would otherwise fail silently
    (→ black screen). request() tries each candidate, so a STALE socket left by a
    crashed/nested comp (newer mtime, nothing listening) falls through to the live one
    instead of hard-failing. mtime is read defensively (a socket can be unlinked between
    glob and stat)."""
    env = os.environ.get("RIKKITISOCK")
    if env:
        return [env]
    rt = os.environ.get("XDG_RUNTIME_DIR") or f"/run/user/{os.getuid()}"
    dated = []
    for p in glob.glob(os.path.join(rt, "rikkiti-ipc-*.sock")):
        try:
            dated.append((os.path.getmtime(p), p))
        except OSError:
            pass  # vanished between glob and stat — skip
    dated.sort(reverse=True)
    return [p for _, p in dated]
SUBSCRIBE = 4
TITLEBAR_H = 30  # mirrors rikkiti-comp.c TITLEBAR_H (SSD titlebar, for --frame)

def _recv_exact(s, n):
    buf = b""
    while len(buf) < n:
        c = s.recv(n - len(buf))
        if not c:
            break
        buf += c
    return buf

def request(typ, payload=b"", retries=3, timeout=3.0):
    """Send one framed message and return the reply payload bytes. Retries the
    whole exchange on a timeout / short reply (the comp can be briefly busy under
    load); raises on a hard failure."""
    candidates = _candidate_socks()
    if not candidates:
        sys.exit("rikkiti-msg: RIKKITISOCK not set and no comp socket found (run inside a Rikkiti session)")
    last = None
    for _ in range(max(1, retries)):
        for sock in candidates:  # a dead newest socket (crashed comp) falls through to the live one
            try:
                s = socket.socket(socket.AF_UNIX)
                s.settimeout(timeout)
                s.connect(sock)
                s.sendall(b"i3-ipc" + struct.pack("=II", len(payload), typ) + payload)
                hdr = _recv_exact(s, 14)
                if len(hdr) < 14:
                    last = f"{sock}: short header"; s.close(); continue
                ln = struct.unpack("=I", hdr[6:10])[0]
                data = _recv_exact(s, ln)
                s.close()
                if len(data) < ln:
                    last = f"{sock}: short body"; continue
                return data
            except (socket.timeout, OSError) as e:
                last = f"{sock}: {e}"
    sys.exit(f"rikkiti-msg: no reply from compositor ({last})")

def request_fd(typ, payload=b"", timeout=5.0):
    """Like request() but also receives ONE fd via SCM_RIGHTS — the comp's
    capture/thumbnail replies pass a memfd alongside the JSON. Returns
    (reply_bytes, fd); fd is -1 if none was attached. The caller owns the fd."""
    candidates = _candidate_socks()
    if not candidates:
        sys.exit("rikkiti-msg: no comp socket found (run inside a Rikkiti session)")
    fdsize = struct.calcsize("i")
    last = None
    for sock in candidates:
        s = None
        try:
            s = socket.socket(socket.AF_UNIX)
            s.settimeout(timeout)
            s.connect(sock)
            s.sendall(b"i3-ipc" + struct.pack("=II", len(payload), typ) + payload)
            # The fd rides the sendmsg, so the FIRST recvmsg carries it — read a
            # generous chunk (header + small JSON) with an ancillary fd buffer.
            chunk, ancdata, _flags, _addr = s.recvmsg(65536, socket.CMSG_LEN(fdsize))
            passed = -1
            for level, ctype, cdata in ancdata:
                if level == socket.SOL_SOCKET and ctype == socket.SCM_RIGHTS:
                    n = len(cdata) // fdsize
                    fds = struct.unpack("=%di" % n, cdata[:n * fdsize])
                    passed = fds[0]
                    for extra in fds[1:]:
                        os.close(extra)  # only one expected; close any surplus
            buf = bytearray(chunk)
            while len(buf) < 14:
                more = s.recv(14 - len(buf))
                if not more:
                    break
                buf += more
            if len(buf) < 14:
                last = f"{sock}: short header"
                if passed >= 0:
                    os.close(passed)
                s.close()
                continue
            ln = struct.unpack("=I", bytes(buf[6:10]))[0]
            while len(buf) - 14 < ln:
                more = s.recv(ln - (len(buf) - 14))
                if not more:
                    break
                buf += more
            s.close()
            return bytes(buf[14:14 + ln]), passed
        except (socket.timeout, OSError) as e:
            last = f"{sock}: {e}"
            if s is not None:
                try:
                    s.close()
                except OSError:
                    pass
    sys.exit(f"rikkiti-msg: no reply from compositor ({last})")


def write_png_rgba(path, w, h, stride, raw):
    """Write RGBA8 bytes (R,G,B,A per pixel; each row `stride` bytes, first w*4
    used) as a PNG using only stdlib zlib — no image-library dependency, so the
    compositor never has to encode in its event loop."""
    def _chunk(typ, data):
        return (struct.pack(">I", len(data)) + typ + data +
                struct.pack(">I", zlib.crc32(typ + data) & 0xffffffff))
    rowbytes = w * 4
    scan = bytearray()
    for y in range(h):
        off = y * stride
        scan.append(0)                       # filter type 0 (none)
        scan += raw[off:off + rowbytes]
    ihdr = struct.pack(">IIBBBBB", w, h, 8, 6, 0, 0, 0)  # 8-bit, colour type 6 (RGBA)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(_chunk(b"IHDR", ihdr))
        f.write(_chunk(b"IDAT", zlib.compress(bytes(scan), 6)))
        f.write(_chunk(b"IEND", b""))


def _emit_capture(payload, out_path, default_name):
    """Send a CAPTURE request (payload bytes), mmap the returned memfd, write it as
    a PNG to out_path (or `default_name` under /tmp), and print the path. Shared by
    every capture shape — the comp hands back the same raw-RGBA+metadata reply."""
    meta_bytes, fd = request_fd(CAPTURE, payload)
    try:
        meta = json.loads(meta_bytes.decode("utf-8", "replace") or "{}")
    except ValueError:
        meta = {}
    if not meta.get("ok") or fd < 0:
        if fd >= 0:
            os.close(fd)
        sys.exit(f"rikkiti-msg: capture failed ({payload.decode('ascii', 'replace')}: {meta})")
    w, h, stride = int(meta["w"]), int(meta["h"]), int(meta["stride"])
    try:
        mm = mmap.mmap(fd, stride * h, access=mmap.ACCESS_READ)
        raw = mm.read(stride * h)
        mm.close()
    finally:
        os.close(fd)
    write_png_rgba(out_path or default_name, w, h, stride, raw)
    print(out_path or default_name)


def _parse_region(pos):
    """Accept 'x,y wxh' (slurp/grim style) or 'x y w h' → (x, y, w, h) ints."""
    parts = " ".join(pos).replace(",", " ").replace("x", " ").split()
    if len(parts) != 4:
        sys.exit("rikkiti-msg: capture region needs 'x,y wxh' or 'x y w h'")
    try:
        x, y, w, h = (int(p) for p in parts)
    except ValueError:
        sys.exit("rikkiti-msg: region coords must be integers")
    if w <= 0 or h <= 0:
        sys.exit("rikkiti-msg: region width/height must be positive")
    return x, y, w, h


def do_capture(rest):
    """rikkiti-msg capture <window|output|region|all> … — compositor-native,
    reliable, NON-INTRUSIVE capture (docs/38); prints the PNG path.
      capture window <id|app_id|--focused> [--max N] [-o PATH]  window content
      capture output <name|--focused>      [--max N] [-o PATH]  one monitor
      capture region <x,y wxh>|<x y w h>   [--max N] [-o PATH]  a rectangle
      capture all                          [--max N] [-o PATH]  whole desktop"""
    shape = "window"
    if rest and rest[0] in ("window", "output", "region", "all"):
        shape = rest[0]; rest = rest[1:]
    # Peel common flags; the remainder is the shape's positional target.
    out_path, max_dim, pos = None, 0, []
    i = 0
    while i < len(rest):
        a = rest[i]
        if a in ("-o", "--out") and i + 1 < len(rest):
            out_path = rest[i + 1]; i += 2; continue
        if a == "--max" and i + 1 < len(rest):
            try:
                max_dim = max(0, int(rest[i + 1]))
            except ValueError:
                sys.exit("rikkiti-msg: --max needs an integer")
            i += 2; continue
        pos.append(a); i += 1

    if shape == "window":
        if not pos:
            sys.exit("rikkiti-msg: capture window needs <id>|<app_id>|--focused")
        target = pos[0]
        if target.isdigit():
            wid = int(target)
        else:
            m = pick(target)
            if not m:
                sys.exit(f"rikkiti-msg: no window matching {target!r}")
            wid = int(m[0].get("id", 0))
        _emit_capture(f"win {wid} {max_dim}".encode(), out_path,
                      f"/tmp/rikkiti-capture-{wid}.png")
    elif shape == "output":
        name = pos[0] if pos else "--focused"
        name = "@focused" if name in ("--focused", "focused") else name
        _emit_capture(f"out {name} {max_dim}".encode(), out_path,
                      "/tmp/rikkiti-capture-output.png")
    elif shape == "region":
        x, y, w, h = _parse_region(pos)
        _emit_capture(f"reg {x} {y} {w} {h} {max_dim}".encode(), out_path,
                      "/tmp/rikkiti-capture-region.png")
    elif shape == "all":
        _emit_capture(f"all {max_dim}".encode(), out_path,
                      "/tmp/rikkiti-capture-all.png")


def windows():
    return json.loads(request(QUERIES["get-windows"]).decode("utf-8", "replace"))

def pick(target):
    """Return the windows matching `target`: '--focused'/'focused' → the focused
    one; otherwise an exact app_id (or title) match. Empty list if none."""
    ws = windows()
    if target in ("--focused", "focused"):
        return [w for w in ws if w.get("focused")]
    return ([w for w in ws if w.get("app_id") == target] or
            [w for w in ws if w.get("title") == target] or
            [w for w in ws if target.lower() in (w.get("title") or "").lower()])

def main():
    args = sys.argv[1:]
    if not args:
        sys.exit("usage: rikkiti-msg <command…> | get-outputs|get-workspaces|get-windows | "
                 "get-window <app_id>|--focused | geom <app_id>|--focused [--frame] | "
                 "wait-window <app_id> [secs] | "
                 "capture window|output|region|all [target] [--max N] [-o PATH] | "
                 "seat click-in|move-in|drag-in <app_id> …")

    if args[0] in ("get-window", "geom"):
        if len(args) < 2:
            sys.exit(f"rikkiti-msg: {args[0]} needs <app_id> or --focused")
        matches = pick(args[1])
        if not matches:
            sys.exit(f"rikkiti-msg: no window matching {args[1]!r}")
        if args[0] == "get-window":
            print(json.dumps(matches[0] if len(matches) == 1 else matches, indent=1))
            return
        # geom → a grim geometry "x,y WxH"
        frame = "--frame" in args[2:]
        r = matches[0].get("rect") or {}
        x, y, w, h = r.get("x", 0), r.get("y", 0), r.get("width", 0), r.get("height", 0)
        if frame and matches[0].get("ssd"):
            y -= TITLEBAR_H; h += TITLEBAR_H  # the SSD titlebar sits above the content rect
        print(f"{x},{y} {w}x{h}")
        return

    if args[0] == "capture":
        do_capture(args[1:])
        return

    if args[0] == "wait-window":
        if len(args) < 2:
            sys.exit("rikkiti-msg: wait-window needs <app_id> or --focused")
        timeout = float(args[2]) if len(args) > 2 else 10.0
        deadline = time.monotonic() + timeout
        while True:
            m = pick(args[1])
            if m:
                print(json.dumps(m[0].get("rect") or {}))
                return
            if time.monotonic() >= deadline:
                sys.exit(f"rikkiti-msg: timed out after {timeout:g}s waiting for {args[1]!r}")
            time.sleep(0.2)

    # Window-relative pointer input: resolve <app_id> → its rect origin and offset
    # the coordinate PAIRS, then forward the absolute `seat <op>` command. Mirrors
    # the absolute grammar (button leads click/dblclick, trails drag):
    #   seat click-in|dblclick-in <app_id> [button] <rx> <ry>
    #   seat move-in              <app_id> <rx> <ry>
    #   seat drag-in              <app_id> <rx1> <ry1> <rx2> <ry2> [button]
    SEAT_IN = {"click-in": "click", "dblclick-in": "dblclick", "move-in": "move", "drag-in": "drag"}
    if args[0] == "seat" and len(args) >= 2 and args[1] in SEAT_IN:
        op = SEAT_IN[args[1]]
        rest = args[2:]
        if not rest:
            sys.exit(f"rikkiti-msg: seat {args[1]} needs <app_id> and coordinates")
        matches = pick(rest[0])
        if not matches:
            sys.exit(f"rikkiti-msg: no window matching {rest[0]!r}")
        r = matches[0].get("rect") or {}
        ox, oy = int(r.get("x", 0)), int(r.get("y", 0))
        rest = rest[1:]
        BTN = ("left", "right", "middle", "back", "forward")
        lead, trail = [], []
        if op in ("click", "dblclick") and rest and rest[0] in BTN:
            lead = [rest[0]]; rest = rest[1:]
        if op == "drag" and rest and rest[-1] in BTN:
            trail = [rest[-1]]; rest = rest[:-1]
        try:
            nums = [float(x) for x in rest]
        except ValueError:
            sys.exit(f"rikkiti-msg: seat {args[1]} coordinates must be numbers")
        if not nums or len(nums) % 2 != 0:
            sys.exit(f"rikkiti-msg: seat {args[1]} needs x y coordinate pair(s)")
        absxy = []
        for i in range(0, len(nums), 2):
            absxy += [str(round(nums[i] + ox)), str(round(nums[i + 1] + oy))]
        args = ["seat", op] + lead + absxy + trail  # forward the absolute command

    if args[0] in QUERIES:
        sys.stdout.write(request(QUERIES[args[0]]).decode("utf-8", "replace").rstrip("\n") + "\n")
        return

    # Anything else is a compositor COMMAND (type 0). The comp replies {"success":…}.
    reply = request(0, " ".join(args).encode())
    sys.stdout.write(reply.decode("utf-8", "replace").rstrip("\n") + "\n")

if __name__ == "__main__":
    main()
