#!/usr/bin/env python3
"""A still's reader, differentially: wai-pack's PNG reader against pack_verify.py's, which
reads the zlib stream with Python's zlib module (extensions/packaging.md §5.7).

The inputs, seeded:

- streams the zlib library writes, at every level and strategy and at windows from
  512 bytes to 32 KiB, and the reference packager's own stills;
- dynamic blocks built here, each code set chosen to sit on one of RFC 1951's and zlib's
  rules: complete codes; no distance code at all; a lone distance code of one bit (the
  one incomplete set zlib takes), used and misused; a lone distance code of 2 to 15 bits;
  over-subscribed and incomplete literal/length, distance and code-length codes; a
  repeat with no length before it, or past the code lengths; no code for the end of the
  block; too many length or distance symbols; a distance past the bytes written; stored
  blocks with a wrong complement; fixed blocks with the symbols they cannot use; and
  several blocks of mixed kinds;
- and every one of those altered: bits flipped, cut short, bytes appended, the header's
  bytes and the Adler-32 changed, chunks dropped, repeated, moved or added, the IHDR's
  fields and a chunk's length changed.

Each one's verdict and, read, its samples' SHA-256 must be the same from both readers,
and the constructions whose verdict the rules decide must have it.

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

Exit status 0 only if every input agrees.
"""
import hashlib
import heapq
import json
import os
import random
import struct
import subprocess
import sys
import zlib

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

SIG = b"\x89PNG\r\n\x1a\n"
ORDER = [16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15]


def chunk(kind, data):
    return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data))


def png_of(w, h, z, extra=()):
    ihdr = struct.pack(">II", w, h) + bytes([8, 2, 0, 0, 0])
    return SIG + chunk(b"IHDR", ihdr) + b"".join(chunk(k, d) for k, d in extra) + chunk(b"IDAT", z) + chunk(b"IEND", b"")


def picture(rng):
    """A picture's filtered rows: random, flat or striped samples, each row with a filter."""
    w, h = rng.randint(1, 40), rng.randint(1, 12)
    kind = rng.randrange(3)
    raw = bytearray()
    for y in range(h):
        raw.append(rng.randrange(5))
        for x in range(3 * w):
            raw.append(rng.randrange(256) if kind == 0 else (7 if kind == 1 else (x // 3 * 11 + y) % 256))
    return w, h, bytes(raw)


class Bits:
    def __init__(self):
        self.out, self.acc, self.n = bytearray(), 0, 0

    def put(self, v, n):
        for i in range(n):
            self.acc |= ((v >> i) & 1) << self.n
            self.n += 1
            if self.n == 8:
                self.out.append(self.acc)
                self.acc, self.n = 0, 0

    def code(self, c, n):
        for i in reversed(range(n)):
            self.put((c >> i) & 1, 1)

    def align(self):
        if self.n:
            self.out.append(self.acc)
            self.acc, self.n = 0, 0


def canonical(lengths):
    """RFC 1951 §3.2.2's codes of `lengths` (an over-subscribed set gives codes anyway)."""
    count = [0] * 16
    for x in lengths:
        count[x] += 1
    count[0] = 0
    nxt, code = [0] * 16, 0
    for b in range(1, 16):
        code = (code + count[b - 1]) << 1
        nxt[b] = code
    out = []
    for x in lengths:
        out.append(nxt[x] if x else 0)
        if x:
            nxt[x] += 1
    return out


def huffman(freq, n, limit):
    """Code lengths for symbols 0..n-1 from their counts (a symbol of count 0 gets none):
    Huffman's, or equal lengths where Huffman's pass `limit`. Complete for two or more
    symbols; one symbol gets one bit."""
    used = [s for s in range(n) if freq.get(s)]
    lengths = [0] * n
    if len(used) == 1:
        lengths[used[0]] = 1
        return lengths
    heap = [(freq[s], i, [s]) for i, s in enumerate(used)]
    heapq.heapify(heap)
    k = len(heap)
    while len(heap) > 1:
        a, b = heapq.heappop(heap), heapq.heappop(heap)
        for s in a[2] + b[2]:
            lengths[s] += 1
        heapq.heappush(heap, (a[0] + b[0], k, a[2] + b[2]))
        k += 1
    if max(lengths) > limit:
        # Equal lengths, then the shortest codes split off until the set is complete.
        bits = max(1, (len(used) - 1).bit_length())
        lengths = [bits if s in used else 0 for s in range(n)]
        spare = (1 << bits) - len(used)
        for s in used:
            if spare == 0:
                break
            lengths[s] -= 1
            spare -= 1
    return lengths


LEN_BASE, LEN_EXTRA = PV.LEN_BASE, PV.LEN_EXTRA
DIST_BASE, DIST_EXTRA = PV.DIST_BASE, PV.DIST_EXTRA


def tokens(raw, dists, rng):
    """`raw` as literals and matches, each match at a distance of a symbol in `dists`."""
    out, i = [], 0
    while i < len(raw):
        best = None
        if dists and rng.random() < 0.8:
            for s in dists:
                for d in {DIST_BASE[s], DIST_BASE[s] + (1 << DIST_EXTRA[s]) - 1}:
                    if d <= i:
                        m = 0
                        while m < 258 and i + m < len(raw) and raw[i + m - d] == raw[i + m]:
                            m += 1
                        if m >= 3 and (best is None or m > best[0]):
                            best = (m, d, s)
        if best:
            out.append(("m", best[0], best[1], best[2]))
            i += best[0]
        else:
            out.append(("l", raw[i]))
            i += 1
    return out


def length_symbol(m):
    li = max(k for k, x in enumerate(LEN_BASE) if x <= m)
    return 257 + li, m - LEN_BASE[li], LEN_EXTRA[li]


def write_tokens(b, toks, lit, lc, dist, dc, bad_dist_bit=False):
    misused = False
    for t in toks:
        if t[0] == "l":
            b.code(lc[t[1]], lit[t[1]])
            continue
        _, m, d, s = t
        sym, extra, nbits = length_symbol(m)
        b.code(lc[sym], lit[sym])
        b.put(extra, nbits)
        if bad_dist_bit and not misused:
            b.put(1, 1)  # the unassigned half of a lone one-bit code
            misused = True
            continue
        b.code(dc[s], dist[s])
        b.put(d - DIST_BASE[s], DIST_EXTRA[s])
    b.code(lc[256], lit[256])


def over(lengths):
    """A complete set made over-subscribed: a code for a symbol that had none, or, with
    every symbol coded, its longest code a bit shorter."""
    free = [x for x in range(len(lengths)) if not lengths[x]]
    if free:
        lengths[free[0]] = 1
    else:
        s = max(range(len(lengths)), key=lambda x: lengths[x])
        lengths[s] -= 1


def dynamic_block(b, toks, last, variant, rng):
    """One dynamic block of `toks` with the code sets `variant` asks for."""
    freq = {256: 1}
    dfreq = {}
    for t in toks:
        s = t[1] if t[0] == "l" else length_symbol(t[1])[0]
        freq[s] = freq.get(s, 0) + 1
        if t[0] == "m":
            dfreq[t[3]] = dfreq.get(t[3], 0) + 1
    if len(freq) == 1:
        freq[0] = 1
    nlen = max(257, max(freq) + 1)
    lit = huffman(freq, nlen, 15)
    if variant == "dist-none":
        ndist, dist = 1, [0]
    elif dfreq:
        ndist = max(dfreq) + 1
        dist = huffman(dfreq, ndist, 15)
    else:
        ndist, dist = 1, [0]
    used_d = [s for s in range(ndist) if dist[s]]
    if variant.startswith("dist-lone-") and len(used_d) == 1:
        dist[used_d[0]] = 1 if variant in ("dist-lone-1", "dist-lone-1-misused") else rng.randint(2, 15)
    elif variant == "dist-over":
        dist = dist + [1, 1, 1]
        ndist = len(dist)
    elif variant == "dist-incomplete" and used_d:
        # One code a bit longer: the set falls short of complete, its longest past one bit.
        s = min(used_d, key=lambda x: dist[x])
        dist[s] += 1
    elif variant == "lit-incomplete":
        s = min((x for x in range(nlen) if lit[x]), key=lambda x: lit[x])
        lit[s] += 1
    elif variant == "lit-over":
        # A code for a symbol that had none, on a set that was complete.
        over(lit)
    elif variant == "no-eob":
        lit[256] = 0
    ndist = min(ndist, 30)
    dist = dist[:ndist]
    lc, dc = canonical(lit), canonical(dist)
    seq = lit + dist
    # The code lengths, run-length coded at random (§3.2.7).
    syms, i = [], 0
    if variant == "rep16-first":
        syms.append((16, 0, 2))
    while i < len(seq):
        x = seq[i]
        run = 1
        while i + run < len(seq) and seq[i + run] == x:
            run += 1
        if x == 0 and run >= 11 and rng.random() < 0.7:
            n = min(run, 138)
            syms.append((18, n - 11, 7))
            i += n
        elif x == 0 and run >= 3 and rng.random() < 0.7:
            n = min(run, 10)
            syms.append((17, n - 3, 3))
            i += n
        elif i > 0 and seq[i - 1] == x and run >= 3 and rng.random() < 0.7:
            n = min(run, 6)
            syms.append((16, n - 3, 2))
            i += n
        else:
            syms.append((x, 0, 0))
            i += 1
    if variant == "rep-past":
        syms.append((17, 7, 3))
    cfreq = {}
    for s, _, _ in syms:
        cfreq[s] = cfreq.get(s, 0) + 1
    if len(cfreq) == 1:
        cfreq[0 if 0 not in cfreq else 1] = 1
    cl = huffman(cfreq, 19, 7)
    if variant == "cl-incomplete":
        s = min((x for x in range(19) if 0 < cl[x] < 7), key=lambda x: cl[x])
        cl[s] += 1
    elif variant == "cl-over":
        over(cl)
    cc = canonical(cl)
    ncode = max(4, max(k for k, s in enumerate(ORDER) if cl[s]) + 1)
    b.put(last, 1)
    b.put(2, 2)
    hlit = nlen - 257 if variant != "hlit-big" else rng.choice([30, 31])
    hdist = ndist - 1 if variant != "hdist-big" else rng.choice([30, 31])
    b.put(hlit, 5)
    b.put(hdist, 5)
    b.put(ncode - 4, 4)
    for s in ORDER[:ncode]:
        b.put(cl[s], 3)
    for s, extra, nbits in syms:
        b.code(cc[s], cl[s])
        b.put(extra, nbits)
    write_tokens(b, toks, lit, lc, dist, dc, bad_dist_bit=variant == "dist-lone-1-misused")


def fixed_block(b, toks, last, variant):
    lit = [8] * 144 + [9] * 112 + [7] * 24 + [8] * 8
    lc = canonical(lit)
    dist, dc = [5] * 32, canonical([5] * 32)
    b.put(last, 1)
    b.put(1, 2)
    if variant == "fixed-lit-286":
        b.code(lc[286], 8)
    elif variant == "fixed-dist-30" and any(t[0] == "m" for t in toks):
        k = next(i for i, t in enumerate(toks) if t[0] == "m")
        toks = toks[:k] + [("m", toks[k][1], toks[k][2], 30)] + toks[k + 1:]
        dist_base = DIST_BASE + [0, 0]
        for t in toks:
            if t[0] == "l":
                b.code(lc[t[1]], 8 if t[1] < 144 else 9)
            else:
                sym, extra, nbits = length_symbol(t[1])
                b.code(lc[sym], lit[sym])
                b.put(extra, nbits)
                b.code(dc[t[3]], 5)
                if t[3] < 30:
                    b.put(t[2] - dist_base[t[3]], DIST_EXTRA[t[3]])
        b.code(lc[256], 7)
        return
    write_tokens(b, toks, lit, lc, dist, dc)


def stored_block(b, data, last, variant):
    b.put(last, 1)
    b.put(0, 2)
    b.align()
    n = len(data)
    nn = (~n) & 0xFFFF if variant != "stored-complement" else n
    b.out += struct.pack("<HH", n, nn) + data


# Each construction, and the verdict RFC 1951 and zlib's inflate give it (None: the
# readers must only agree).
VARIANTS = {
    "complete": "ok", "dist-none": "ok", "dist-lone-1": "ok", "dist-lone-1-misused": "refused",
    "dist-lone-L": "refused", "dist-over": "refused", "dist-incomplete": "refused",
    "lit-incomplete": "refused", "lit-over": "refused", "cl-incomplete": "refused", "cl-over": "refused",
    "rep16-first": "refused", "rep-past": "refused", "no-eob": "refused", "hlit-big": "refused",
    "hdist-big": "refused", "far": "refused", "stored": "ok", "stored-complement": "refused",
    "fixed-lit-286": "refused", "fixed-dist-30": "refused", "mixed": "ok",
}


def construct(rng, variant):
    w, h, raw = picture(rng)
    b = Bits()
    if variant in ("stored", "stored-complement"):
        stored_block(b, raw, 1, variant)
    elif variant in ("fixed-lit-286", "fixed-dist-30"):
        toks = tokens(raw, list(range(4)), rng)
        if variant == "fixed-dist-30" and not any(t[0] == "m" for t in toks):
            return None
        fixed_block(b, toks, 1, variant)
    elif variant == "mixed":
        cut = sorted(rng.sample(range(len(raw) + 1), 2)) if len(raw) >= 2 else [0, len(raw)]
        parts = [raw[:cut[0]], raw[cut[0]:cut[1]], raw[cut[1]:]]
        for k, part in enumerate(parts):
            last = int(k == 2)
            kind = rng.randrange(3)
            if kind == 0:
                stored_block(b, part, last, "stored")
            else:
                # Matches only within this part, so a distance never reaches into another.
                toks = tokens(part, list(range(4)), rng)
                if kind == 1:
                    fixed_block(b, toks, last, "")
                else:
                    dynamic_block(b, toks, last, "complete", rng)
    else:
        dists = []
        if variant not in ("dist-none",):
            dists = [rng.randrange(4)] if variant.startswith("dist-lone-") else rng.sample(range(8), rng.randint(1, 4))
        toks = tokens(raw, dists, rng)
        if variant.startswith("dist-lone-") or variant in ("dist-over", "dist-incomplete"):
            if not any(t[0] == "m" for t in toks):
                return None
        if variant == "far":
            toks = [("m", 3, 1, 0)] + toks
        dynamic_block(b, toks, 1, "dist-lone-L" if variant == "dist-lone-L" else variant, rng)
    b.align()
    z = b"\x78\x01" + bytes(b.out) + struct.pack(">I", zlib.adler32(raw))
    return png_of(w, h, z)


def zlib_streams(rng):
    out = []
    for _ in range(40):
        w, h, raw = picture(rng)
        co = zlib.compressobj(rng.choice([0, 1, 6, 9]), zlib.DEFLATED, rng.choice([9, 10, 12, 15]), 9,
                              rng.choice([zlib.Z_DEFAULT_STRATEGY, zlib.Z_FIXED, zlib.Z_HUFFMAN_ONLY, zlib.Z_RLE]))
        out.append(png_of(w, h, co.compress(raw) + co.flush()))
        rgb = bytes(rng.randrange(256) if rng.random() < 0.5 else 9 for _ in range(3 * w * h))
        out.append(PV.png_still(rgb, w, h))
    return out


def chunks_of(b):
    at, out = 8, []
    while at + 12 <= len(b):
        n = struct.unpack_from(">I", b, at)[0]
        if at + 12 + n > len(b):
            break
        out.append((b[at + 4:at + 8], b[at + 8:at + 8 + n]))
        at += 12 + n
    return out


def build(chs):
    return SIG + b"".join(chunk(k, d) for k, d in chs)


def mutate(b, rng):
    if len(b) < 2:
        return b + bytes([rng.randrange(256)])
    chs = chunks_of(b)
    op = rng.randrange(12)
    if op == 0 or not chs:
        b = bytearray(b)
        for _ in range(rng.randint(1, 3)):
            i = rng.randrange(len(b))
            b[i] ^= 1 << rng.randrange(8)
        return bytes(b)
    if op == 1:
        return b[:rng.randrange(len(b))]
    if op == 2:
        i, j = rng.randrange(len(chs)), rng.randrange(len(chs))
        chs[i], chs[j] = chs[j], chs[i]
        return build(chs)
    if op == 3:
        chs.insert(rng.randrange(len(chs) + 1), chs[rng.randrange(len(chs))])
        return build(chs)
    if op == 4:
        del chs[rng.randrange(len(chs))]
        return build(chs)
    if op == 5:
        kind = rng.choice([b"tEXt", b"PLTE", b"IDAT", b"IEND", b"IHDR", b"abCD", b"ABCD", b"gAMA"])
        data = bytes(rng.randrange(256) for _ in range(rng.choice([0, 1, 3, 6, 13, 768, 771])))
        chs.insert(rng.randrange(len(chs) + 1), (kind, data))
        return build(chs)
    idats = [i for i, (k, _) in enumerate(chs) if k == b"IDAT"]
    if op in (6, 7, 8, 9) and idats:
        z = bytearray(b"".join(chs[i][1] for i in idats))
        if op == 6 and z:
            for _ in range(rng.randint(1, 4)):
                i = rng.randrange(len(z))
                z[i] ^= 1 << rng.randrange(8)
        elif op == 7:
            z = z[:rng.randrange(len(z) + 1)]
        elif op == 8:
            z += bytes(rng.randrange(256) for _ in range(rng.randint(1, 6)))
        elif op == 9 and len(z) >= 2:
            z[rng.randrange(2)] = rng.randrange(256)
            if rng.random() < 0.7:
                flg = z[1] & 0xE0
                z[1] = flg | (31 - (z[0] * 256 + flg) % 31) % 31
        z = bytes(z)
        cut = rng.randrange(len(z) + 1)
        rest = [c for i, c in enumerate(chs) if i not in idats]
        pos = idats[0]
        return build(rest[:pos] + [(b"IDAT", z[:cut]), (b"IDAT", z[cut:])] + rest[pos:])
    if op == 10:
        for i, (k, d) in enumerate(chs):
            if k == b"IHDR" and len(d) == 13:
                d = bytearray(d)
                d[rng.randrange(13)] = rng.choice([0, 1, 2, 3, 6, 8, 16, 255])
                chs[i] = (k, bytes(d))
        return build(chs)
    return b + bytes(rng.randrange(256) for _ in range(rng.randint(1, 12)))


def py_verdict(b):
    try:
        w, h, px = PV.png_frame(b)
        return {"verdict": "ok", "w": w, "h": h, "rgb_sha256": hashlib.sha256(px).hexdigest()}
    except ValueError:
        return {"verdict": "refused"}


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 2000
    rng = random.Random(seed)
    cases = []  # (bytes, variant or None)
    for f in sorted(os.listdir(os.path.join(HERE, "stills"))):
        if f.endswith(".png"):
            cases.append((open(os.path.join(HERE, "stills", f), "rb").read(), None))
    cases += [(b, None) for b in zlib_streams(rng)]
    names = sorted(VARIANTS)
    while len(cases) < count // 2:
        v = names[len(cases) % len(names)]
        b = construct(rng, v)
        if b is not None:
            cases.append((b, v))
    base = [b for b, _ in cases]
    while len(cases) < count:
        b = rng.choice(base)
        for _ in range(rng.choice([1, 1, 2, 3])):
            b = mutate(b, rng)
        cases.append((b, None))
    lines = "".join(json.dumps({"op": "read-png", "hex": b.hex()}) + "\n" for b, _ in cases)
    out = subprocess.run([binary, "stage-json"], input=lines.encode(), capture_output=True, check=True).stdout.decode().splitlines()
    if len(out) != len(cases):
        print(f"FAIL wai-pack answered {len(out)} of {len(cases)} lines")
        return 1
    fails, seen = 0, {"ok": 0, "refused": 0}
    by_variant = {}
    for (b, v), got in zip(cases, out):
        rust, mine = json.loads(got), py_verdict(b)
        seen[mine["verdict"]] += 1
        if v is not None:
            by_variant[v] = by_variant.get(v, 0) + 1
        why = None
        if rust != mine:
            why = f"wai-pack {json.dumps(rust)}, pack_verify.py {json.dumps(mine)}"
        elif v is not None and VARIANTS[v] != mine["verdict"]:
            why = f"the construction {v} should be {VARIANTS[v]}, both readers say {mine['verdict']}"
        if why:
            fails += 1
            if fails <= 5:
                print(f"FAIL {v or 'input'} {b[:64].hex()}…: {why}")
    missing = [v for v in names if v not in by_variant]
    if missing:
        fails += 1
        print(f"FAIL constructions never built: {missing}")
    print(f"png_diff seed {seed}: {len(cases) - fails}/{len(cases)} agree ({seen['ok']} read, {seen['refused']} refused; "
          f"{sum(by_variant.values())} constructions of {len(by_variant)} kinds)")
    return 0 if fails == 0 else 1


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