#!/usr/bin/env python3
"""The source readers, differentially: wai-pack against pack_verify.py on seeded mutations.

Generates Y4M and WAV files from a fixed seed: headers with every parameter the readers
treat (dimensions at and past their range, frame rates, interlacing, colour spaces,
FRAME-line parameters), WAV format chunks with every tag, channel count (2^15 and
2^16 − 1 among them), sample size, block alignment and chunk length the readers treat,
chunks in either order, and truncations and single-byte changes of each. Every file's
verdict from `wai-pack source-verdict` must equal pack_verify.py's, and no reader may
panic (`PANIC`).

    python3 source_diff.py <wai-pack binary> [seed] [count]

Exit status 0 only if every verdict agrees. The run is reproducible: Python's `random`
with a fixed seed.
"""
import os
import random
import struct
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 y4m(rng):
    w, h = rng.randint(1, 5), rng.randint(1, 5)
    toks = [b"YUV4MPEG2", b"W%d" % w, b"H%d" % h]
    c = rng.choice([b"", b" C420jpeg", b" C444", b" C420", b" C422", b" C420p10", b" Cmono"])
    extra = rng.choice([b"", b" F25:1", b" F0:1", b" F25", b" Ip", b" It", b" I", b" A1:1", b" XYSCSS=420JPEG", b"  ",
                        b" W0", b" H99999999999", b" W4294967295", b" F4294967296:1"])
    hdr = b" ".join(toks) + c + extra + b"\n"
    cw, ch = (w, h) if c == b" C444" else ((w + 1) // 2, (h + 1) // 2)
    size = w * h + 2 * cw * ch
    body = b""
    for _ in range(rng.randint(0, 3)):
        fl = rng.choice([b"FRAME\n", b"FRAME\n", b"FRAME Ip\n", b"FRAME Ib\n", b"FRAME W4\n", b"FRAME X1\n", b"FRAMEX\n",
                         b"FRAME  \n", b"FRAME I\n"])
        body += fl + bytes(rng.randrange(256) for _ in range(size))
    b = hdr + body
    r = rng.random()
    if r < 0.3 and b:
        b = b[:rng.randrange(len(b))]
    elif r < 0.5 and b:
        i = rng.randrange(len(b))
        b = b[:i] + bytes([rng.randrange(256)]) + b[i + 1:]
    return b


def wav(rng):
    tag = rng.choice([1, 1, 1, 3, 0xFFFE])
    ch = rng.choice([1, 2, 6, 0, 32768, 32769, 65535])
    bits = rng.choice([16, 16, 8, 24])
    align = rng.choice([2 * ch & 0xFFFF, 2 * ch & 0xFFFF, 0, 4])
    rate = rng.choice([8000, 0])
    fl = rng.choice([16, 18, 40, 41, 15, 17])
    fmt = struct.pack("<HHIIHH", tag, ch, rate, rate * 2 & 0xFFFFFFFF, align, bits) + bytes(max(fl - 16, 0))
    fmt = fmt[:fl]
    chunks = []
    if rng.random() < 0.3:
        chunks.append(b"LIST" + struct.pack("<I", 3) + b"abc" + (b"\0" if rng.random() < 0.7 else b""))
    chunks.append(b"fmt " + struct.pack("<I", fl) + fmt + (b"\0" if fl % 2 else b""))
    nd = rng.randint(0, 3) * 2 * max(ch, 1) if ch < 100 else rng.choice([0, 2 * ch])
    dl = rng.choice([nd, nd, nd + 1, 0xFFFFFFFE])
    chunks.append(b"data" + struct.pack("<I", dl & 0xFFFFFFFF) + bytes(nd))
    if rng.random() < 0.2:
        chunks.reverse()
    b = b"RIFF\0\0\0\0WAVE" + b"".join(chunks)
    if rng.random() < 0.25:
        b = b[:rng.randrange(len(b))]
    return b


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 2
    binary = argv[1]
    seed = int(argv[2]) if len(argv) > 2 else 1
    count = int(argv[3]) if len(argv) > 3 else 3000
    rng = random.Random(seed)
    with tempfile.TemporaryDirectory() as d:
        paths = []
        for i in range(count):
            p = os.path.join(d, f"{i}.y4m" if i % 2 else f"{i}.wav")
            with open(p, "wb") as f:
                f.write(y4m(rng) if i % 2 else wav(rng))
            paths.append(p)
        rust = {}
        for k in range(0, len(paths), 500):
            r = subprocess.run([binary, "source-verdict"] + paths[k:k + 500], capture_output=True)
            if r.returncode != 0:
                print(r.stderr.decode(errors="replace")[-2000:])
                return 1
            for line in r.stdout.decode().splitlines():
                a, v = line.split("\t")
                rust[a] = v
        bad = 0
        for p in paths:
            with open(p, "rb") as f:
                b = f.read()
            py = PV.source_verdict(p, b)
            if rust.get(p) != py:
                bad += 1
                if bad <= 10:
                    print(f"DIFF {os.path.basename(p)}: wai-pack {rust.get(p)}, pack_verify.py {py}")
    print(f"source_diff: seed {seed}, {count} files, {bad} differences")
    return 0 if bad == 0 else 1


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