#!/usr/bin/env python3
"""Write a golden package's supplied floor (extensions/packaging.md §7.3) from its fixture.

The floor of a golden is an AV1 coding of its generated fixture, so its every frame comes
from the fixture generator and nothing else. This script states how it was made:

1. `wai-pack fixture gen <name>` writes the fixture the policy's component names;
2. `wai-pack plan` cuts it with the policy and writes the keyframe plan (§7.3);
3. for each floor rung of the policy, in ladder order, each piece is coded alone by an AV1
   encoder (`aomenc`, one pass, constant quality, no lag and no alternate reference, so
   each temporal unit holds one shown frame, and no key frame but the piece's first), in
   the low-overhead OBU format; a rung given `--scale <rung>=<d>` codes the fixture with
   its width and height divided by d, each sample the mean of its d × d block, rounded
   half up, and `--cq <rung>=<n>` sets its quality (40 by default);
4. each rung's pieces' streams, concatenated in order, are `floor/<floor name>.obu`: the
   component's id when it has one floor rung, the rung's id when it has several (§7.1);
5. the policy's `floor.supplied_listing` is set to the SHA-256 of the floor directory's
   listing (SPEC §3.1), which `wai-pack listing` prints.

An encoder is not a pure function of its input across versions, so the stream is committed
and pinned by its listing: a packer checks the floor against the listing and the plan, and
never encodes it. This script is how the committed stream was written, not a step of
verification.

    python3 supplied_floor.py <wai-pack binary> <packs/<name> directory> [aomenc]
                              [--scale <rung>=<d>]... [--cq <rung>=<n>]...
"""
import os
import re
import subprocess
import sys
import tempfile


def run(*args):
    return subprocess.run(args, check=True, capture_output=True).stdout


def frames_of(y4m):
    """The stream header line and each frame's bytes (its FRAME line included)."""
    head_end = y4m.index(b"\n") + 1
    head = y4m[:head_end]
    w = int(re.search(rb" W(\d+)", head).group(1))
    h = int(re.search(rb" H(\d+)", head).group(1))
    size = w * h + 2 * ((w + 1) // 2) * ((h + 1) // 2)
    out, at = [], head_end
    while at < len(y4m):
        line = y4m.index(b"\n", at) + 1
        out.append(y4m[at:line + size])
        at = line + size
    return head, out


def scaled(head, frames, d):
    """The frames with width and height divided by d: each sample of each plane the mean
    of its d x d block, rounded half up (4:2:0, so the chroma planes are halved too)."""
    if d == 1:
        return head, frames
    w = int(re.search(rb" W(\d+)", head).group(1))
    h = int(re.search(rb" H(\d+)", head).group(1))
    if w % (2 * d) or h % (2 * d):
        raise SystemExit(f"{w}x{h} does not divide by {2 * d}")
    cw, ch = (w + 1) // 2, (h + 1) // 2

    def plane(b, pw, ph):
        out = bytearray()
        for y in range(0, ph, d):
            for x in range(0, pw, d):
                t = sum(b[(y + j) * pw + x + i] for j in range(d) for i in range(d))
                out.append((2 * t + d * d) // (2 * d * d))
        return bytes(out)
    out = []
    for f in frames:
        line_end = f.index(b"\n") + 1
        body = f[line_end:]
        y, u, v = body[:w * h], body[w * h:w * h + cw * ch], body[w * h + cw * ch:]
        out.append(f[:line_end] + plane(y, w, h) + plane(u, cw, ch) + plane(v, cw, ch))
    head = re.sub(rb" W\d+", b" W%d" % (w // d), head)
    head = re.sub(rb" H\d+", b" H%d" % (h // d), head)
    return head, out


def floor_rungs(policy):
    """The policy's floor rungs, in ladder order."""
    out = []
    for block in re.split(r"\n(?=\[)", policy):
        if block.startswith("[[ladder]]") and re.search(r'role = "floor"', block):
            out.append(re.search(r'id = "([^"]+)"', block).group(1))
    return out


def main(argv):
    opts = {"--scale": {}, "--cq": {}}
    rest = []
    it = iter(argv[1:])
    for a in it:
        if a in opts:
            k, v = next(it).split("=")
            opts[a][k] = int(v)
        else:
            rest.append(a)
    if len(rest) < 2:
        print(__doc__)
        return 2
    wai_pack, pack = rest[0], os.path.abspath(rest[1])
    aomenc = rest[2] if len(rest) > 2 else "aomenc"
    policy_path = os.path.join(pack, "policy.toml")
    policy = open(policy_path, encoding="utf-8").read()
    comp = re.search(r'\[\[component\]\]\s*id = "([^"]+)"', policy).group(1)
    source = re.search(r'source = "([^"]+)\.y4m"', policy).group(1)
    rungs = floor_rungs(policy)
    names = [comp] if len(rungs) == 1 else rungs
    floor = os.path.join(pack, "floor")
    with tempfile.TemporaryDirectory() as tmp:
        src = os.path.join(tmp, "source.y4m")
        run(wai_pack, "fixture", "gen", source, "-o", src)
        plan = os.path.join(tmp, "plan.txt")
        run(wai_pack, "plan", "--policy", policy_path, "--component", comp, "--source", src, "--out", plan)
        starts = [int(x) for x in open(plan).read().split()]
        head0, frames0 = frames_of(open(src, "rb").read())
        bounds = starts + [len(frames0)]
        streams = {}
        for rung, name in zip(rungs, names):
            head, frames = scaled(head0, frames0, opts["--scale"].get(rung, 1))
            stream = b""
            for k, (a, b) in enumerate(zip(bounds, bounds[1:])):
                piece = os.path.join(tmp, f"{rung}-p{k}.y4m")
                with open(piece, "wb") as f:
                    f.write(head + b"".join(frames[a:b]))
                out = os.path.join(tmp, f"{rung}-p{k}.obu")
                run(aomenc, "--obu", f"--limit={b - a}", "--passes=1", "--lag-in-frames=0", "--auto-alt-ref=0",
                    "--kf-min-dist=0", f"--kf-max-dist={b - a}", "--end-usage=q", f"--cq-level={opts['--cq'].get(rung, 40)}",
                    "--cpu-used=6", "--threads=1", "--row-mt=0", "-q", "-o", out, piece)
                stream += open(out, "rb").read()
            streams[name] = stream
        os.makedirs(floor, exist_ok=True)
        for n in os.listdir(floor):
            os.remove(os.path.join(floor, n))
        for name, stream in streams.items():
            with open(os.path.join(floor, f"{name}.obu"), "wb") as f:
                f.write(stream)
    import hashlib
    listing = hashlib.sha256(run(wai_pack, "listing", floor)).hexdigest()
    policy = re.sub(r'supplied_listing = "[0-9a-f]{64}"', f'supplied_listing = "{listing}"', policy)
    open(policy_path, "w", encoding="utf-8").write(policy)
    for name, stream in streams.items():
        print(f"{floor}/{name}.obu: {len(stream)} bytes, {len(frames0)} frames, key frames at {starts}")
    print(f"listing {listing}")
    return 0


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