#!/usr/bin/env python3
"""The policy reader, differentially: wai-pack against pack_verify.py.

Two sets of policies, each file's verdict (`ok`, or the reason a packager refuses it)
and, when the policy itself parses, its digest, from `wai-pack policy-verdict`, must
equal pack_verify.py's:

- byte-level cases: a byte-order mark, bytes that are not UTF-8, integers at and past
  2^53 and 2^63, nesting at and past the cap in arrays, inline tables and dotted keys,
  control characters, line endings, NUL, escapes, ids at the byte limit, one id named
  twice across kinds;
- seeded random policies with random faults in every table (from a fixed seed).

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

Exit status 0 only if every verdict and digest agrees. Python 3.11 or later (tomllib).
"""
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

MIN = '''schema = "wai.pack.policy/1"
asset_id = "a"
lead = "v"
[[component]]
id = "v"
media = "video"
timebase = { num = 25, den = 1 }
[[ladder]]
id = "b0"
component = "v"
capability = "wai.video.int_motion"
role = "base"
'''


def byte_cases():
    c = {}

    def case(name, b):
        c[name] = b.encode() if isinstance(b, str) else b
    case("min", MIN)
    case("bom", b"\xef\xbb\xbf" + MIN.encode())
    case("int-2^53", MIN + "zzz = 9007199254740992\n")
    case("int-2^63", MIN + "zzz = 9223372036854775808\n")
    case("int-2^64", MIN + "zzz = 18446744073709551616\n")
    case("int-i64max", MIN + "zzz = 9223372036854775807\n")
    case("int-i64min", MIN + "zzz = -9223372036854775808\n")
    case("hex-int", MIN + "[segmentation]\ntarget_ms = 0x7d0\n")
    case("nan", MIN + "zzz = nan\n")
    case("neg-zero", MIN + "[segmentation]\ncut_window = -0\n")
    case("plus", MIN + "[segmentation]\ncut_window = +8\n")
    case("non-utf8", MIN.encode() + b"# \xff\xfe comment\n")
    case("non-utf8-str", MIN.replace('asset_id = "a"', 'asset_id = "a\xff"').encode("latin-1"))
    case("surrogate-utf8", MIN.encode() + b"# \xed\xa0\x80\n")
    for d in (31, 32, 33, 60, 79, 80, 200, 3000):
        case(f"deep-array-{d}", "zzz = " + "[" * d + "]" * d + "\n" + MIN)
    for d in (31, 32, 200):
        case(f"deep-inline-{d}", "zzz = " + "{a=" * d + "1" + "}" * d + "\n" + MIN)
    for d in (31, 32, 2000):
        case(f"deep-dotted-{d}", "zzz." + ".".join(["a"] * d) + " = 1\n" + MIN)
    case("deep-after-float", "aaa = 1.5\nzzz = " + "[" * 40 + "]" * 40 + "\n" + MIN)
    case("dup-key", MIN + "asset_id = \"b\"\n")
    case("ctrl-del", MIN.replace('asset_id = "a"', 'asset_id = "a\\u007f\\u001f\\b\\f"'))
    case("u2028", MIN.replace('asset_id = "a"', 'asset_id = "a\\u2028\\u2029"'))
    case("astral", MIN.replace('asset_id = "a"', 'asset_id = "\\U0001F600"'))
    case("id-255-bytes", MIN.replace('asset_id = "a"', 'asset_id = "' + "é" * 127 + 'a"'))
    case("id-256-bytes", MIN.replace('asset_id = "a"', 'asset_id = "' + "é" * 128 + '"'))
    case("crlf", MIN.replace("\n", "\r\n"))
    case("lone-cr", MIN.replace("\n", "\r"))
    case("same-id-across-kinds", MIN.replace('id = "b0"', 'id = "v"'))
    case("tab-in-literal", MIN.replace('asset_id = "a"', "asset_id = 'a\tb'"))
    case("float-exp", MIN + "zzz = 1e3\n")
    case("local-time", MIN + "zzz = 07:32:00\n")
    case("int-refine-base", MIN.replace("wai.video.int_motion", "wai.video.int_refine"))
    case("empty", b"")
    case("nul", MIN.encode() + b"\x00")
    case("u0000-escape", MIN.replace('asset_id = "a"', 'asset_id = "a\\u0000"'))
    case("key-unicode", MIN + '"é" = 1\n')
    case("ladder-empty", MIN.split("[[ladder]]")[0] + "ladder = []\n")
    case("mixed-array", MIN + 'zzz = [1, "a"]\n')
    case("params-wrong-type", MIN + "params = 5\n")
    case("cut-window-2", MIN + "[segmentation]\ncut_window = 2\n")
    return c


CAPS = ["wai.video.int_motion", "wai.video.int_refine", "wai.video.av1", "wai.audio.flac", "wai.image.png",
        "wai.audio.int_refine", "wai.audio.int_layered", "wai.splat.sog", "wai.video.avc", "wai.video.package",
        "wai.meta.hdr10plus", "wai.nope"]
PARAMS = ['{}', '{ block = 8 }', '{ block = 5 }', '{ search = 33 }', '{ q = 6 }', '{ rate = 8000 }',
          '{ q_candidates = [] }', '{ q_candidates = [16, 12] }', '{ source = "encode" }', '{ bogus = 1 }',
          '{ to_rate = 48000, q = 64 }']
KINDS = ['"static"', '"bogus"', '']


def random_policy(rng):
    def biased(seq):
        seq = list(seq)
        return seq[0] if (rng.random() < 0.85 and seq) else rng.choice(seq)

    def val(kind):
        if rng.random() < 0.01:
            return biased(['1.5', '-3', '9007199254740992', '"x"', 'true', '[]', '{}'])
        return kind

    q = '"'
    L = []
    if rng.random() > 0.03:
        L.append("schema = " + val(q + "wai.pack.policy/1" + q))
    if rng.random() > 0.03:
        L.append("asset_id = " + val(q + biased(["a", "", "x" * 256]) + q))
    if rng.random() < 0.5:
        L.append("royalty_free_only = " + val(biased(["true", "false"])))
    comps = [rng.choice(["v", "a", "v2", "v"]) for _ in range(biased([1, 2, 0, 3]))]
    lead = rng.choice(comps if rng.random() < 0.95 else comps + ["zz"]) if comps else "v"
    L.append(f'lead = "{lead}"')
    if rng.random() < 0.05:
        L.append("zz_unknown = 1")
    body = []
    for c in comps:
        body += ["[[component]]", f'id = "{c}"', "media = " + val(q + biased(["video", "audio", "bogus"]) + q)]
        # §7.8's profile and a source, which a floor-only component may not have.
        if rng.random() < 0.15:
            body.append("profile = " + val(q + biased(["floor-only", "bogus"]) + q))
        if rng.random() < 0.1:
            body.append('source = "v.y4m"')
        if rng.random() > 0.05:
            num, den = val(str(biased([25, 0, 48000, 4294967295, 16777216]))), val(str(biased([1, 1001, 0, 600])))
            body.append("timebase = { num = " + num + ", den = " + den + " }")
    if rng.random() < 0.6:
        body.append("[segmentation]")
        for k, choices in [("target_ms", [2000, 500, 0]), ("min_ms", [1000, 3000, 0]), ("max_ms", [4000, 100, 600000, 600001, 349525]),
                           ("cut_abs_permille", [300, 0, 1001]), ("cut_window", [8, 0, 2, 3, 65]),
                           ("lookahead_frames", [6, 65]), ("closed", ["true", "false"])]:
            if rng.random() < 0.4:
                body.append(f"{k} = " + val(str(biased(choices))))
    rungs = []
    for i in range(biased([2, 1, 3, 4, 0])):
        rid = biased([f"r{i}", "r0", "b", "v"])
        role = rng.choice(["base", "refinement", "base", "floor", "bogus"])
        rungs.append(rid)
        comp = rng.choice(comps if rng.random() < 0.95 else comps + ["zz"]) if comps else "v"
        cap = rng.choice(CAPS[:5] if rng.random() < 0.8 else CAPS)
        body += ["[[ladder]]", f'id = "{rid}"', f'component = "{comp}"', f'capability = "{cap}"', f'role = "{role}"']
        if role == "refinement" or rng.random() < 0.1:
            body.append(f'refines = "{rng.choice(rungs + ["nope"])}"')
        if rng.random() < 0.5:
            body.append("params = " + biased(PARAMS))
    if rng.random() < 0.3:
        # §7.7's form of a supplied floor, and the floor's other members.
        body.append("[floor]")
        if rng.random() < 0.6:
            body.append("supplied_form = " + val(q + biased(["cmaf", "obu"]) + q))
        if rng.random() < 0.4:
            body.append("extra_keyframes = " + val(biased(["true", "false"])))
    if rng.random() < 0.3:
        body += ["[live]", "layer_lag_ms = " + biased(["[0]", "[0, 200]", "[0, 200, 400]", "[1]", "[0, 100, 50]"])]
    for i in range(biased([0, 0, 1, 2])):
        body += ["[[class]]", f'id = "{biased(["c" + str(i), "c0"])}"']
        if rng.random() < 0.5:
            body.append(f"max_layer = {biased([0, 1, 2, 3])}")
        if rng.random() < 0.3:
            body.append(f"max_kbps = {biased([1, 0, 900])}")
    for i in range(biased([0, 0, 1, 2])):
        body += ["[[moment]]", f'id = "m{i}"', f"from_ms = {biased([0, 100])}", f"to_ms = {biased([0, 100, 200])}"]
        if rng.random() < 0.5:
            body.append('tags = ["event"]')
    for i in range(biased([0, 0, 1, 2, 5])):
        body += ["[[rule]]", f'id = "{biased(["u" + str(i), "u0"])}"']
        if rng.random() < 0.5:
            body.append("when = { piece_kind = [" + biased(KINDS) + "] }")
        t = []
        if rng.random() < 0.4:
            pick = rng.sample(rungs + ["nope"], k=min(2, len(rungs) + 1))
            t.append("ladder = [" + ", ".join(q + x + q for x in pick) + "]")
        if rng.random() < 0.4:
            t.append(f"max_layer = {biased([0, 1, 2])}")
        if rng.random() < 0.4:
            t.append(f"deadline_ms = {biased([1, 100, 200, 300, 400, 800])}")
        if t:
            body.append("then = { " + ", ".join(t) + " }")
    if rng.random() < 0.3:
        body += ["[defaults]", f"deadline_ms = {biased([1500, 1, 0, 150])}"]
    if rng.random() < 0.2:
        body += ["[classify]", f"synthetic_bins = {biased([12, 65])}"]
    return ("\n".join(L + body) + "\n").encode()


def python_verdict(b):
    try:
        P = PV.parse_policy_bytes(b)
    except PV.Refused as e:
        return e.code, "-"
    try:
        PV.admit(b)
        v = "ok"
    except PV.Refused as e:
        v = e.code
    return v, PV.digest(P)


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)
    cases = list(byte_cases().items()) + [(f"random-{i}", random_policy(rng)) for i in range(count)]
    bad = 0
    with tempfile.TemporaryDirectory() as d:
        paths = {}
        for name, b in cases:
            p = os.path.join(d, name + ".toml")
            with open(p, "wb") as f:
                f.write(b)
            paths[p] = (name, b)
        rust = {}
        files = list(paths)
        for k in range(0, len(files), 500):
            r = subprocess.run([binary, "policy-verdict"] + files[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():
                p, v, dg = line.split("\t")
                rust[p] = (v, dg)
        for p, (name, b) in paths.items():
            py = python_verdict(b)
            if rust.get(p) != py:
                bad += 1
                if bad <= 10:
                    print(f"DIFF {name}: wai-pack {rust.get(p)}, pack_verify.py {py}")
    print(f"policy_diff: seed {seed}, {len(cases)} policies, {bad} differences")
    return 0 if bad == 0 else 1


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