#!/usr/bin/env python3
"""Verify the hand-rolled FlatBuffers encoder against the reference runtime.

The C code in native/src/net/flatbuf.c builds Hyperion protocol messages
without flatcc. This decodes those bytes using the upstream `flatbuffers`
Python package, so a layout mistake fails here rather than silently producing
a message HyperHDR drops on the floor.

Schema (hyperion.ng libsrc/flatbufserver/hyperion_request.fbs):

    table Register { origin:string (required); priority:int; }
    table RawImage { data:[ubyte]; width:int = -1; height:int = -1; }
    table Image   { data:ImageType (required); duration:int = -1; }
    table Clear   { priority:int; }
    union ImageType { RawImage, NV12Image }        // RawImage = 1
    union Command   { Color, Image, Clear, Register }  // Image = 2, Register = 4
    table Request { command:Command (required); }
    root_type Request;
"""

import struct
import subprocess
import sys
import tempfile
from pathlib import Path

from flatbuffers import number_types as N
from flatbuffers.table import Table

CMD_IMAGE = 2
CMD_REGISTER = 4
IMGTYPE_RAWIMAGE = 1

FAILURES = []


def check(label, actual, expected):
    ok = actual == expected
    status = "ok  " if ok else "FAIL"
    shown = actual if not isinstance(actual, (bytes, bytearray)) else bytes(actual).hex()
    exp = expected if not isinstance(expected, (bytes, bytearray)) else bytes(expected).hex()
    print(f"  [{status}] {label}: {shown!r}" + ("" if ok else f" (expected {exp!r})"))
    if not ok:
        FAILURES.append(label)


def unframe(raw: bytes) -> bytes:
    """Strip and validate the 4-byte big-endian length prefix."""
    assert len(raw) >= 4, "message shorter than its length prefix"
    (declared,) = struct.unpack(">I", raw[:4])
    check("length prefix matches payload", declared, len(raw) - 4)
    return raw[4:]


def root_table(payload: bytes) -> Table:
    pos = struct.unpack_from("<I", payload, 0)[0]
    return Table(bytearray(payload), pos)


def field(tbl: Table, slot: int):
    """Return the vtable offset for a slot, or 0 when the field is absent."""
    return tbl.Offset(slot * 2 + 4)


def read_u8(tbl: Table, slot: int, default=0):
    o = field(tbl, slot)
    return tbl.Get(N.Uint8Flags, o + tbl.Pos) if o else default


def read_i32(tbl: Table, slot: int, default=0):
    o = field(tbl, slot)
    return tbl.Get(N.Int32Flags, o + tbl.Pos) if o else default


def read_sub(tbl: Table, slot: int):
    o = field(tbl, slot)
    if not o:
        return None
    return Table(tbl.Bytes, tbl.Indirect(o + tbl.Pos))


def read_str(tbl: Table, slot: int):
    o = field(tbl, slot)
    return tbl.String(o + tbl.Pos).decode() if o else None


def read_bytes(tbl: Table, slot: int):
    o = field(tbl, slot)
    if not o:
        return None
    start = tbl.Vector(o)
    length = tbl.VectorLen(o)
    return bytes(tbl.Bytes[start : start + length])


def verify_register(raw: bytes):
    print("Register message:")
    req = root_table(unframe(raw))
    check("Request.command_type", read_u8(req, 0), CMD_REGISTER)

    reg = read_sub(req, 1)
    assert reg is not None, "Request.command missing"
    check("Register.origin", read_str(reg, 0), "lgtv-audio-cap")
    check("Register.priority", read_i32(reg, 1), 150)


def verify_image(raw: bytes):
    print("Image message:")
    req = root_table(unframe(raw))
    check("Request.command_type", read_u8(req, 0), CMD_IMAGE)

    img = read_sub(req, 1)
    assert img is not None, "Request.command missing"
    check("Image.data_type", read_u8(img, 0), IMGTYPE_RAWIMAGE)
    # duration defaults to -1 and is omitted from the buffer.
    check("Image.duration (default)", read_i32(img, 2, default=-1), -1)

    raw_img = read_sub(img, 1)
    assert raw_img is not None, "Image.data missing"
    check("RawImage.width", read_i32(raw_img, 1, default=-1), 4)
    check("RawImage.height", read_i32(raw_img, 2, default=-1), 2)

    expected = bytes(b for i in range(8) for b in (i * 10, i * 10 + 1, i * 10 + 2))
    check("RawImage.data length", len(read_bytes(raw_img, 0) or b""), 24)
    check("RawImage.data contents", read_bytes(raw_img, 0), expected)


def main():
    repo = Path(__file__).resolve().parent.parent
    sources = [
        repo / "test" / "fb_dump.c",
        repo / "native" / "src" / "net" / "hyperion.c",
        repo / "native" / "src" / "net" / "flatbuf.c",
        repo / "native" / "src" / "common" / "log.c",
    ]

    with tempfile.TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        binary = tmp / "fb_dump"

        compile_cmd = ["cc", "-std=c11", "-Wall", "-Wextra", "-O1", "-o", str(binary)]
        compile_cmd += [str(s) for s in sources]
        print("$ " + " ".join(compile_cmd))
        subprocess.run(compile_cmd, check=True)

        for kind, verifier in (("register", verify_register), ("image", verify_image)):
            out = tmp / f"{kind}.bin"
            subprocess.run([str(binary), kind, str(out)], check=True)
            verifier(out.read_bytes())

    if FAILURES:
        print(f"\n{len(FAILURES)} check(s) failed: {', '.join(FAILURES)}")
        return 1
    print("\nAll FlatBuffers checks passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())