#!/usr/bin/env python3
"""
End-to-end smoke test for krane-fb-stub under qemu-system-aarch64.

- test.dtb: krane-sku176.dtb with a /firmware/coreboot node added (as
  depthcharge's fixup produces), reg pair 1 -> synthetic LBIO table at
  0x45000000, written into guest RAM by a qemu loader device.
- The LBIO table describes a 64x16 xrgb framebuffer at 0x45200000.
- The stub paints red -> yellow -> green -> blue, ~2s apart, into that buffer.
  We sample the framebuffer color via QEMU monitor pmemsave snapshots and
  verify the full checkpoint sequence plus final hold-on-blue.

Usage: python3 qemu_test.py
"""
import os
import select
import socket
import struct
import subprocess
import time

LBIO_ADDR = 0x45000000
FB_ADDR = 0x45200000
FB_X, FB_Y, FB_BPL = 64, 16, 256

RED, YELLOW, GREEN, BLUE = 0x00FF0000, 0x00FFFF00, 0x0000FF00, 0x000000FF


def build_lbio_table():
    t = bytearray(24 + 48)
    t[0:4] = b"LBIO"
    t[4:8] = struct.pack("<I", 24)      # header_bytes
    t[20:24] = struct.pack("<I", 1)     # table_entries
    t[24:28] = struct.pack("<I", 0x12)  # LB_TAG_FRAMEBUFFER
    t[28:32] = struct.pack("<I", 40)    # record size
    t[32:40] = struct.pack("<Q", FB_ADDR)
    t[40:44] = struct.pack("<I", FB_X)
    t[44:48] = struct.pack("<I", FB_Y)
    t[48:52] = struct.pack("<I", FB_BPL)
    t[52] = 32                          # bits_per_pixel
    t[53], t[54] = 16, 8                # red
    t[55], t[56] = 8, 8                 # green
    t[57], t[58] = 0, 8                 # blue
    return bytes(t)


def make_test_dtb():
    subprocess.run(["cp", "krane-sku176.dtb", "test.dtb"], check=True)
    subprocess.run(["dtc", "-I", "dtb", "-O", "dts", "test.dtb", "-o", "test.dts"],
                   check=True)
    with open("test.dts", "a") as f:
        f.write("""
&{/} {
	firmware {
		#address-cells = <0x02>;
		#size-cells = <0x02>;
		ranges;

		coreboot {
			compatible = "coreboot";
			reg = <0x00 0x45000000 0x00 0x380
			       0x00 0x45100000 0x00 0x127000>;
		};
	};
};
""")
    subprocess.run(["dtc", "-I", "dts", "-O", "dtb", "test.dts", "-o", "test.dtb"],
                   check=True)
    out = subprocess.run(["dtc", "-I", "dtb", "-O", "dts", "test.dtb"],
                         capture_output=True, text=True).stdout
    assert "reg = <0x00 0x45000000 0x00 0x380" in out, "coreboot reg missing"
    print("test.dtb ready (with /firmware/coreboot)")


def main():
    make_test_dtb()
    with open("lbio.bin", "wb") as f:
        f.write(build_lbio_table())

    parent, child = socket.socketpair()
    qemu = subprocess.Popen(
        ["qemu-system-aarch64", "-M", "virt", "-m", "1024",
         "-nographic", "-serial", "null", "-monitor", "stdio", "-cpu", "max",
         "-kernel", "krane-fb-stub.bin", "-dtb", "test.dtb",
         "-device", f"loader,file=lbio.bin,addr=0x{LBIO_ADDR:x},force-raw=on"],
        stdin=child.fileno(), stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    child.close()
    time.sleep(2)

    def drain():
        out = b""
        while select.select([qemu.stdout], [], [], 0.05)[0]:
            c = os.read(qemu.stdout.fileno(), 65536)
            if not c:
                break
            out += c
        return out

    drain()

    def pmem(addr, size, path):
        if os.path.exists(path):
            os.remove(path)
        parent.sendall(f'pmemsave 0x{addr:x} {size} "{path}"\n'.encode())
        for _ in range(60):
            drain()
            if os.path.exists(path) and os.path.getsize(path) == size:
                return open(path, "rb").read()
            time.sleep(0.1)
        raise RuntimeError(path)

    fails = 0

    def check(name, ok, detail=""):
        nonlocal fails
        print(f"{name:34s} {'PASS' if ok else 'FAIL'} {detail}")
        if not ok:
            fails += 1

    # Full timeline: sample the first framebuffer pixel every 0.5s for
    # ~9s of stub time and require the exact transition sequence
    # red -> yellow -> green -> blue, each state held, blue never left.
    # qemu startup eats ~1-2s, so sampling starts inside the red window.
    seq = []
    for i in range(22):
        snap = pmem(FB_ADDR, 4, "/tmp/fb_seq.bin")
        c = struct.unpack_from("<I", snap, 0)[0]
        if not seq or seq[-1][1] != c:
            seq.append((i * 0.5, c))
        time.sleep(0.5)

    # Timeline runs 11s wall (~9s stub) and ends well inside blue. So
    # instead of a yellow full-screen snapshot, verify full-screen uniformity
    # only in blue (the final, held state) — the timeline already proves
    # red/yellow/green appeared in order.
    time.sleep(4)
    t9 = pmem(FB_ADDR, FB_X * FB_Y * 4, "/tmp/fb_t9.bin")
    parent.sendall(b"quit\n")
    qemu.wait(timeout=10)

    def px(d, i):
        return struct.unpack_from("<I", d, i * 4)[0]

    N = FB_X * FB_Y
    order = [c for _, c in seq]
    check("sequence red->yellow->green->blue",
          order[:4] == [RED, YELLOW, GREEN, BLUE],
          f"(observed {[(f'{t:.1f}s', f'0x{c:08x}') for t, c in seq]})")
    check("blue held forever (no further transitions)", len(order) == 4,
          f"({len(order)} distinct states)")
    check("t=13s(stub) full screen uniform blue",
          all(px(t9, i) == 0x000000FF for i in (0, N // 2, N - 1)),
          f"(0x{px(t9,0):08x})")
    print("\nALL PASS" if not fails else f"\nFAILED: {fails}")
    raise SystemExit(1 if fails else 0)


if __name__ == "__main__":
    main()