"""Two framing helpers of make_external.sh.

    python3 -I frame.py jxsvideo <out> <w> <h> <fps_num> <fps_den> <codestreams>
        Splits a file of consecutive JPEG XS codestreams (each with its length in its
        picture header's Lcod, as a constant-bit-rate encoder writes) and writes them in
        wai.video.jxs's framing: five u32 little-endian values n, h, w, fps_num, fps_den,
        then each codestream preceded by its u32 little-endian length (format-probes §3.3).

    python3 -I frame.py base <in> <out>
        Keeps the NAL units of an H.265 Annex B byte stream whose nuh_layer_id is 0: the
        base layer of a multilayer bitstream, each NAL unit after a four-byte start code.
"""

import struct
import sys


def codestreams(b):
    out, at = [], 0
    while at < len(b):
        if b[at:at + 2] != b"\xff\x10" or b[at + 2:at + 4] != b"\xff\x50":
            raise SystemExit("no SOC and CAP at %d" % at)
        lcap = struct.unpack(">H", b[at + 4:at + 6])[0]
        pih = at + 4 + lcap
        if b[pih:pih + 2] != b"\xff\x12":
            raise SystemExit("no PIH at %d" % pih)
        lcod = struct.unpack(">I", b[pih + 4:pih + 8])[0]
        if lcod == 0:
            raise SystemExit("a codestream without Lcod: the encoder was not at a constant rate")
        out.append(b[at:at + lcod])
        at += lcod
    return out


def jxsvideo(out, w, h, num, den, src):
    cs = codestreams(open(src, "rb").read())
    data = struct.pack("<5I", len(cs), int(h), int(w), int(num), int(den))
    for c in cs:
        data += struct.pack("<I", len(c)) + c
    open(out, "wb").write(data)


def nal_units(b):
    starts, i = [], 0
    while i + 2 < len(b):
        if b[i] == 0 and b[i + 1] == 0 and b[i + 2] == 1:
            starts.append(i + 3)
            i += 3
        else:
            i += 1
    units = []
    for k, s in enumerate(starts):
        e = starts[k + 1] - 3 if k + 1 < len(starts) else len(b)
        while e > s and b[e - 1] == 0:
            e -= 1
        units.append(b[s:e])
    return units


def base(src, out):
    keep = [u for u in nal_units(open(src, "rb").read()) if ((u[0] << 8 | u[1]) >> 3) & 0x3F == 0]
    open(out, "wb").write(b"".join(b"\x00\x00\x00\x01" + u for u in keep))


if __name__ == "__main__":
    {"jxsvideo": jxsvideo, "base": base}[sys.argv[1]](*sys.argv[2:])