#!/usr/bin/env python3
"""Decode a captured kitty-graphics stream back into images.

We can't eyeball the emulator's output on a terminal that doesn't speak the
protocol, so this reads a raw capture of stdout, reassembles every
transmit-and-place command (including chunked payloads and tmux passthrough
wrapping) and writes each image out as a PPM/PAM, reporting the control keys.

    ./build/sl0pboy --kitty 3 --kitty-b64 rom.gb > cap.bin
    tools/kitty_decode.py cap.bin -o /tmp/frames

Use --check FILE.ppm to assert an image matches a --shot capture pixel for pixel.
"""
import argparse
import base64
import os
import re
import sys

APC = re.compile(rb"\x1b_G(.*?)\x1b\\", re.S)
TMUX = re.compile(rb"\x1bPtmux;(.*?)\x1b\\\Z", re.S)


def unwrap_tmux(data: bytes) -> bytes:
    """Undo tmux passthrough framing: DCS tmux; <ESC-doubled payload> ST."""
    out = bytearray()
    i = 0
    while i < len(data):
        if data.startswith(b"\x1bPtmux;", i):
            j = i + 7
            body = bytearray()
            while j < len(data):
                if data[j] == 0x1B:
                    if j + 1 < len(data) and data[j + 1] == 0x1B:
                        body.append(0x1B)   # doubled ESC -> literal ESC
                        j += 2
                        continue
                    if j + 1 < len(data) and data[j + 1] == ord("\\"):
                        j += 2              # ST: end of passthrough
                        break
                body.append(data[j])
                j += 1
            out += body
            i = j
        else:
            out.append(data[i])
            i += 1
    return bytes(out)


def parse_ctrl(ctrl: bytes):
    keys = {}
    for part in ctrl.split(b","):
        if b"=" in part:
            k, v = part.split(b"=", 1)
            keys[k.decode()] = v.decode()
    return keys


def images(data: bytes):
    """Yield (control_keys, payload_bytes) for each complete transmission."""
    cur = None
    for m in APC.finditer(data):
        body = m.group(1)
        ctrl, _, payload = body.partition(b";")
        keys = parse_ctrl(ctrl)
        if cur is None:
            cur = (keys, bytearray())
        cur[1].extend(payload)
        if keys.get("m", "0") == "0":
            try:
                yield cur[0], base64.b64decode(bytes(cur[1]))
            except Exception as e:      # truncated capture (killed mid-frame)
                print(f"skipping undecodable payload: {e}", file=sys.stderr)
            cur = None


def write_image(path, keys, raw):
    w, h = int(keys["s"]), int(keys["v"])
    fmt = int(keys.get("f", "24"))
    bpp = 4 if fmt == 32 else 3
    want = w * h * bpp
    if len(raw) != want:
        raise SystemExit(f"{path}: payload {len(raw)} bytes, expected {want}")
    if bpp == 3:
        with open(path, "wb") as f:
            f.write(b"P6\n%d %d\n255\n" % (w, h))
            f.write(raw)
    else:                                   # RGBA -> PAM keeps the alpha
        with open(path, "wb") as f:
            f.write(b"P7\nWIDTH %d\nHEIGHT %d\nDEPTH 4\nMAXVAL 255\n"
                    b"TUPLTYPE RGB_ALPHA\nENDHDR\n" % (w, h))
            f.write(raw)
    return w, h, bpp


def read_ppm(path):
    with open(path, "rb") as f:
        blob = f.read()
    if not blob.startswith(b"P6"):
        raise SystemExit(f"{path}: not a P6 PPM")
    fields, pos = [], 2
    while len(fields) < 3:
        while pos < len(blob) and blob[pos : pos + 1].isspace():
            pos += 1
        if blob[pos : pos + 1] == b"#":
            while blob[pos] != 0x0A:
                pos += 1
            continue
        start = pos
        while not blob[pos : pos + 1].isspace():
            pos += 1
        fields.append(int(blob[start:pos]))
    return fields[0], fields[1], blob[pos + 1 :]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("capture")
    ap.add_argument("-o", "--outdir", default=None)
    ap.add_argument("--check", help="PPM the first RGB image must match exactly")
    ap.add_argument("--quiet", action="store_true")
    args = ap.parse_args()

    data = unwrap_tmux(open(args.capture, "rb").read())
    if args.outdir:
        os.makedirs(args.outdir, exist_ok=True)

    n = 0
    checked = False
    for keys, raw in images(data):
        n += 1
        if keys.get("a") == "d":            # delete/cleanup, no pixels
            if not args.quiet:
                print(f"[{n}] delete i={keys.get('i')} d={keys.get('d')}")
            continue
        if keys.get("t") == "s":            # payload is the shm object name
            if not args.quiet:
                print(f"[{n}] shm transfer, object={raw.decode()!r} "
                      f"{keys.get('s')}x{keys.get('v')} i={keys.get('i')}")
            continue
        ext = "pam" if keys.get("f") == "32" else "ppm"
        path = os.path.join(args.outdir or ".", f"img{n:04d}.{ext}")
        if args.outdir:
            w, h, bpp = write_image(path, keys, raw)
        else:
            w, h = int(keys["s"]), int(keys["v"])
            bpp = 4 if keys.get("f") == "32" else 3
        if not args.quiet:
            print(f"[{n}] {w}x{h} bpp={bpp} i={keys.get('i')} p={keys.get('p')} "
                  f"z={keys.get('z','0')} c={keys.get('c','-')} "
                  f"r={keys.get('r','-')} q={keys.get('q')} "
                  f"C={keys.get('C')} bytes={len(raw)}"
                  + (f" -> {path}" if args.outdir else ""))
        if args.check and not checked and bpp == 3:
            cw, ch, pix = read_ppm(args.check)
            if (cw, ch) != (w, h):
                raise SystemExit(f"size mismatch: kitty {w}x{h} vs "
                                 f"{args.check} {cw}x{ch}")
            if pix[: cw * ch * 3] != raw:
                bad = sum(a != b for a, b in zip(pix, raw))
                raise SystemExit(f"pixel mismatch: {bad} bytes differ")
            print(f"OK: image {n} matches {args.check} exactly ({w}x{h})")
            checked = True

    if not n:
        raise SystemExit("no kitty graphics commands found in capture")
    if args.check and not checked:
        raise SystemExit("no RGB image to check")
    print(f"{n} graphics command(s) decoded")


if __name__ == "__main__":
    main()