#!/usr/bin/env python3
"""
err_apply.py - put one error case from the ERR lists into a clean stream (4-Codec Extended Test Sets, ERR group).

  err_apply.py STREAM CASE_ID [CASE_ID ...]        damage STREAM with the given cases (one output file per case)
  err_apply.py STREAM --all [--kind bitflip,burst] [--region header|payload]
                                                   every case of STREAM from the list (optionally filtered)
  err_apply.py STREAM --list                       print the cases of STREAM and exit

  --cases FILE   case list (.csv or .json); default: <av1|vp9|avc|hvc>_err_v1.csv picked from the case id prefix,
                 looked up next to this script, in ./errors/ and in ../errors/
  -o DIR         output directory (default: current directory)
  --force        apply even if STREAM's file name differs from the list's stream column

Output: <stream name>__<case id>.<ext>, e.g. ERR_Long_854x480_..._2.52Mbps__VP9-ERR_Long-hbf-00012.ivf,
and a line per case in DIR/err_apply_manifest.csv (case id, output file, kind, region, byte, bit, length, md5 of the output).

Error kinds (byte offsets in the file, bit 0 = MSB of the byte):
  bitflip   flip bit `bit` of byte `byte`
  burst     overwrite `length` bytes from `byte` with random bytes; Python random.Random(seed).randrange(256) per byte
            (use this script to reproduce them: another RNG gives other bytes)
  truncate  cut the file at `byte`
  drop      remove `length` bytes from `byte`; in an IVF file the size field of the frame(s) it cuts into is reduced
Python 3 standard library only.
"""
import argparse, csv, hashlib, json, os, random, struct, sys

PREFIX = {"AV1": "av1", "VP9": "vp9", "AVC": "avc", "HVC": "hvc"}
INT_FIELDS = ("byte", "bit", "length", "seed")


def apply(data, case):
    """the damaged bytes of `data` for one case (same algorithm as the runner that produced the reference outcomes)"""
    k = case["kind"]; b = case["byte"]; d = bytearray(data)
    if k == "bitflip":
        d[b] ^= 0x80 >> case["bit"]
    elif k == "burst":
        rng = random.Random(case.get("seed") or 0); e = min(b + case["length"], len(d))
        for i in range(b, e): d[i] = rng.randrange(256)
    elif k == "truncate":
        del d[b:]
    elif k == "drop":
        n = case["length"]; fix_ivf_sizes(d, b, n); del d[b:b + n]
    else:
        raise ValueError(f"unknown kind {k}")
    return bytes(d)


def fix_ivf_sizes(d, b, n):
    """shrink the size field of the IVF frame(s) that contain bytes [b, b+n) of the removed range"""
    if d[:4] != b"DKIF": return
    p = 32
    while p + 12 <= len(d):
        sz = struct.unpack_from("<I", d, p)[0]; s, e = p + 12, p + 12 + sz
        cut = max(0, min(e, b + n) - max(s, b))
        if cut: struct.pack_into("<I", d, p, sz - cut)
        p = e


def load_cases(path):
    if path.endswith(".json"):
        rows = json.load(open(path))
    else:
        with open(path, newline="") as f: rows = list(csv.DictReader(f))
    for r in rows:
        for k in INT_FIELDS:
            v = r.get(k)
            r[k] = None if v in (None, "") else int(v)
    return rows


def find_list(case_id):
    codec = PREFIX.get(case_id.split("-")[0])
    if not codec: sys.exit(f"cannot tell the codec from case id {case_id!r}; give the list with --cases")
    here = os.path.dirname(os.path.abspath(__file__))
    for d in (here, os.path.join(os.getcwd(), "errors"), os.path.join(here, "..", "errors"), os.getcwd()):
        for ext in (".csv", ".json"):
            p = os.path.join(d, f"{codec}_err_v1{ext}")
            if os.path.exists(p): return p
    sys.exit(f"{codec}_err_v1.csv not found; download it from the ERR page and give it with --cases")


def main():
    ap = argparse.ArgumentParser(description="Put error cases from the ERR lists into a clean stream.",
                                 formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__)
    ap.add_argument("stream"); ap.add_argument("ids", nargs="*")
    ap.add_argument("--cases"); ap.add_argument("-o", "--out", default=".")
    ap.add_argument("--all", action="store_true"); ap.add_argument("--list", action="store_true")
    ap.add_argument("--kind"); ap.add_argument("--region", choices=("header", "payload"))
    ap.add_argument("--force", action="store_true")
    a = ap.parse_args()
    name = os.path.basename(a.stream); stem, ext = os.path.splitext(name)
    if not a.cases:
        if a.ids: a.cases = find_list(a.ids[0])
        else:   # --all / --list: the codec from the stream's folder or extension
            codec = next((c for c in ("VP9", "AVC", "HVC") if f"/{c}/" in os.path.abspath(a.stream).replace(os.sep, "/")), None)
            codec = codec or {".264": "AVC", ".265": "HVC"}.get(ext)
            if not codec: sys.exit("give the case list with --cases (cannot tell AV1 from VP9 by the file alone)")
            a.cases = find_list(codec + "-x")
    rows = load_cases(a.cases)
    mine = [r for r in rows if os.path.basename(r["stream"]) == name]
    if a.ids:
        by = {r["id"]: r for r in rows}; sel = []
        for i in a.ids:
            if i not in by: sys.exit(f"case {i} not in {a.cases}")
            r = by[i]
            if os.path.basename(r["stream"]) != name and not a.force:
                sys.exit(f"case {i} belongs to {os.path.basename(r['stream'])}, not {name} (--force to apply anyway)")
            sel.append(r)
    elif a.all or a.list:
        if not mine: sys.exit(f"no cases for {name} in {a.cases}")
        sel = [r for r in mine if (not a.kind or r["kind"] in a.kind.split(",")) and (not a.region or r["region"] == a.region)]
    else:
        ap.error("give case ids, --all or --list")
    if a.list:
        for r in sel:
            print(f"{r['id']}\t{r['region']}\t{r['kind']}\tbyte {r['byte']}" + (f" bit {r['bit']}" if r["bit"] is not None else "")
                  + (f" len {r['length']}" if r["length"] else "") + f"\t{r.get('part') or ''} {r.get('field') or ''}\tref: {r.get('ref_outcome') or ''}")
        return
    data = open(a.stream, "rb").read(); os.makedirs(a.out, exist_ok=True)
    man = os.path.join(a.out, "err_apply_manifest.csv"); new = not os.path.exists(man)
    with open(man, "a", newline="") as mf:
        w = csv.writer(mf)
        if new: w.writerow(["id", "file", "kind", "region", "byte", "bit", "length", "md5"])
        for r in sel:
            if r["byte"] is None or r["byte"] > len(data): sys.exit(f"case {r['id']}: byte {r['byte']} outside the stream ({len(data)} bytes) - wrong stream?")
            dmg = apply(data, r); out = os.path.join(a.out, f"{stem}__{r['id']}{ext}")
            open(out, "wb").write(dmg); md5 = hashlib.md5(dmg).hexdigest()
            w.writerow([r["id"], os.path.basename(out), r["kind"], r["region"], r["byte"], r["bit"], r["length"], md5])
            print(f"{r['id']}: {r['region']} {r['kind']} at byte {r['byte']}" + (f" bit {r['bit']}" if r["bit"] is not None else "")
                  + (f" len {r['length']}" if r["length"] else "") + f" -> {out}")


if __name__ == "__main__":
    main()