#!/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. 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;
4. the pieces' streams, concatenated in order, are `floor/<component>.obu`;
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]
"""
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 main(argv):
    if len(argv) < 3:
        print(__doc__)
        return 2
    wai_pack, pack = argv[1], os.path.abspath(argv[2])
    aomenc = argv[3] if len(argv) > 3 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)
    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()]
        head, frames = frames_of(open(src, "rb").read())
        bounds = starts + [len(frames)]
        stream = b""
        for k, (a, b) in enumerate(zip(bounds, bounds[1:])):
            piece = os.path.join(tmp, f"p{k}.y4m")
            with open(piece, "wb") as f:
                f.write(head + b"".join(frames[a:b]))
            out = os.path.join(tmp, f"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", "--cq-level=40", "--cpu-used=6",
                "--threads=1", "--row-mt=0", "-q", "-o", out, piece)
            stream += open(out, "rb").read()
        floor = os.path.join(pack, "floor")
        os.makedirs(floor, exist_ok=True)
        for name in os.listdir(floor):
            os.remove(os.path.join(floor, name))
        with open(os.path.join(floor, f"{comp}.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)
    print(f"{floor}/{comp}.obu: {len(stream)} bytes, {len(frames)} frames, key frames at {starts}; listing {listing}")
    return 0


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