#!/usr/bin/env python3
"""The supplied CMAF reader differentially (extensions/packaging.md §7.7): wai-pack's
(`wai-pack cmaf-verdict`) against pack_verify.py's, which shares no code with it, on
seeded alterations of real files. The files are the CMAF golden's floor (a media muxer's
initialization and media segments, packs/video-cmaf-01/floor/) and §7.6's writer's (the
ladder golden's package, packs/video-ladder-01/package/floor/). Each case alters the
initialization segment or one media segment: a byte flipped in its first boxes, a 32-bit
field overwritten at a box boundary, the file cut short, a box's size changed, a box
inserted at the top level or in a moof, the init's timescale or edit list altered, or a
segment's moof and mdat repeated. Every verdict (ok, floor-cmaf-malformed or
floor-cmaf-timing), and for an accepted segment its samples (offset, length, duration,
flags, decode time), must agree, and no reader may panic.

    python3 cmaf_diff.py <wai-pack binary> <seed> <cases>
"""
import hashlib
import os
import random
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import pack_verify as PV  # noqa: E402


def py_init(b):
    try:
        return "ok", PV.cmaf_init(b)
    except PV.CmafFault as x:
        return x.code, None
    except Exception as x:  # a reader must refuse, never fail otherwise
        return f"CRASH {type(x).__name__}", None


def py_segment(b, init):
    try:
        s = PV.cmaf_segment(b, init)
    except PV.CmafFault as x:
        return (x.code, 0, "-")
    except Exception as x:
        return (f"CRASH {type(x).__name__}", 0, "-")
    text = "".join(f"{a},{n},{d},{f},{t};" for a, n, d, f, t in s)
    return ("ok", len(s), hashlib.sha256(text.encode()).hexdigest())


def boxes(b, start, end):
    try:
        return PV.bmff_boxes(b, start, end)
    except PV.CmafFault:
        return []


def alter(rng, b):
    b = bytearray(b)
    top = boxes(b, 0, len(b))
    k = rng.randrange(8)
    if k == 6:
        # The times: an init's timescale or its edit list's media time, at their offsets
        # from the box type. `mdhd`'s timescale follows its version and flags and its two
        # times, 4 bytes each in version 0 and 8 in version 1.
        for kind, off in ((b"mdhd", None), (b"elst", 20), (b"elst", 16)):
            at = bytes(b).find(kind)
            if at >= 0 and off is None:
                off = 16 if at + 4 < len(b) and b[at + 4] == 0 else 24
            if at >= 0 and rng.random() < 0.5 and at + off + 4 <= len(b):
                b[at + off:at + off + 4] = rng.choice([0, 1, 2, rng.randrange(1 << 32)]).to_bytes(4, "big")
                return bytes(b)
        k = 0
    if k == 7:
        # A media segment's moof and mdat repeated: the second's decode time is the first's.
        pairs = [i for i, x in enumerate(top) if x[0] == b"moof" and i + 1 < len(top) and top[i + 1][0] == b"mdat"]
        if pairs:
            i = rng.choice(pairs)
            b += b[top[i][1]:top[i + 1][3]]
            return bytes(b)
        k = 0
    if k == 0 and b:
        b[rng.randrange(min(len(b), 400))] ^= 1 << rng.randrange(8)
    elif k == 1 and top:
        x = rng.choice(top)
        inner = boxes(b, x[2], x[3]) if x[0] in (b"moov", b"moof", b"trak", b"traf", b"mdia", b"minf", b"stbl", b"mvex") else []
        y = rng.choice(inner) if inner and rng.random() < 0.7 else x
        at = y[1] + rng.choice([8, 12, 16, 20, 24])
        if at + 4 <= len(b):
            b[at:at + 4] = rng.choice([0, 1, 7, 0xFFFFFFFF, rng.randrange(1 << 32)]).to_bytes(4, "big")
    elif k == 2 and len(b) > 8:
        del b[rng.randrange(8, len(b)):]
    elif k == 3 and top:
        x = rng.choice(top)
        b[x[1]:x[1] + 4] = max(0, (x[3] - x[1]) + rng.choice([-9, -1, 1, 4, 8, -(x[3] - x[1])])).to_bytes(4, "big")
    elif k == 4 and top:
        kind = rng.choice([b"free", b"skip", b"uuid", b"styp", b"mdat", b"moof"])
        box = (8 + 4).to_bytes(4, "big") + kind + bytes(4)
        at = rng.choice([x[1] for x in top] + [len(b)])
        b[at:at] = box
    elif top:
        moofs = [x for x in top if x[0] == b"moof"]
        if moofs:
            x = rng.choice(moofs)
            inner = boxes(b, x[2], x[3])
            at = rng.choice([y[1] for y in inner] + [x[3]])
            kind = rng.choice([b"sbgp", b"free", b"tfdt", b"trun"])
            box = (8 + 8).to_bytes(4, "big") + kind + bytes(8)
            b[at:at] = box
            b[x[1]:x[1] + 4] = (x[3] - x[1] + len(box)).to_bytes(4, "big")
    return bytes(b)


def main(argv):
    if len(argv) != 4:
        print(__doc__)
        return 2
    wai_pack, seed, n = argv[1], int(argv[2]), int(argv[3])
    rng = random.Random(seed)
    root = os.path.join(HERE, "packs")
    sets = []
    for d in (os.path.join(root, "video-cmaf-01", "floor", "hi"), os.path.join(root, "video-cmaf-01", "floor", "lo"),
              os.path.join(root, "video-ladder-01", "package", "floor", "hi")):
        names = sorted(os.listdir(d))
        sets.append((PV.load(os.path.join(d, "init.mp4")), [PV.load(os.path.join(d, x)) for x in names if x.endswith(".m4s")]))
    bad, seen = 0, {}
    with tempfile.TemporaryDirectory() as tmp:
        for c in range(n):
            init, segs = rng.choice(sets)
            seg = rng.choice(segs)
            if rng.random() < 0.25:
                init = alter(rng, init)
            else:
                seg = alter(rng, seg)
            ip, sp = os.path.join(tmp, "init.mp4"), os.path.join(tmp, "seg.m4s")
            open(ip, "wb").write(init)
            open(sp, "wb").write(seg)
            out = subprocess.run([wai_pack, "cmaf-verdict", ip, sp], capture_output=True, text=True).stdout.splitlines()
            iv, pi = py_init(init)
            want = [f"init\t{iv}"]
            if iv == "ok":
                v = py_segment(seg, pi)
                want.append(f"{sp}\t{v[0]}\t{v[1]}\t{v[2]}")
            verdict = want[-1].split("\t")[1]
            seen[verdict] = seen.get(verdict, 0) + 1
            if out != want or any("PANIC" in x or "CRASH" in x for x in out + want):
                bad += 1
                if bad <= 10:
                    print(f"case {c}: wai-pack {out}, pack_verify {want}")
    print(f"cmaf_diff: {n - bad}/{n} cases agree; verdicts {dict(sorted(seen.items()))}")
    return 0 if bad == 0 else 1


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