#!/usr/bin/env python3
"""Write a golden's supplied Opus floor (extensions/packaging.md §7.9) in the CMAF form.

The golden is multi/<name>/: its audio component's source is the stereo WAV of the formula
multi_verify.py and wai-pack's multi_fixture write (never committed), at 48 kHz, spanning the
lead's plan. ffmpeg's libopus encodes it at 128 kb/s in packets of 20 ms, and ffmpeg's HLS
muxer writes it as fMP4 segments of 0.2 s, cut on the decode timeline. Each piece's media
segment is then the muxer's fragments of the piece, as they are: the first fragment's `styp`
and every fragment's `moof` and `mdat`, in order (a `sidx` describes one fragment, so none is
kept); each `moof` addresses its samples from its own first byte, so the fragments are
unchanged. The fragments that decode at or past the programme's end, the packet that carries
its last 6.5 ms after the pre-skip, are the floor's tail (§7.9). `init.mp4` is the muxer's.
The files are `multi/<name>/floor/<f>/`, `<f>` the component's id, and the policy's
`floor.supplied_listing` is set to the SHA-256 of the listing (SPEC §3.1) of the golden's
whole supplied floor: the lead's `packs/<lead>/floor/` with this directory beside it.

An encoder and a muxer are not pure functions of their input across versions, so the files
are committed and pinned by the listing: a packer reads them and checks them against the
plan, and never writes them. This script is how they were written, not a step of
verification.

    python3 supplied_opus.py <ffmpeg> <pack-conformance dir> <golden name> <lead pack> <component id>

A floor whose pieces are all 2 s needs no regrouping: ffmpeg's own segments
(`-hls_time 2`) are one per piece, and a tail where the programme ends on a piece boundary.
"""
import hashlib
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile

RATE = 48000
FPS = 25
STEP = 0.2


def wav_sample(k, c):
    if c == 0:
        return ((k * 7) % 400) * 80 - 16000
    t = k % 200
    return (t * 300 if t < 100 else (200 - t) * 300) - 15000


def wav(frames):
    data = struct.pack("<%dh" % (2 * frames), *[wav_sample(k, c) for k in range(frames) for c in (0, 1)])
    return (b"RIFF" + struct.pack("<I", 36 + len(data)) + b"WAVEfmt " + struct.pack("<IHHIIHH", 16, 1, 2, RATE, RATE * 4, 4, 16)
            + b"data" + struct.pack("<I", len(data)) + data)


def boxes(b):
    at, out = 0, []
    while at < len(b):
        size, kind = struct.unpack(">I4s", b[at:at + 8])
        out.append((kind, b[at:at + size]))
        at += size
    return out


def tfdt(seg):
    for kind, body in boxes(seg):
        if kind == b"moof":
            i = body.find(b"tfdt")
            ver = body[i + 4]
            return int.from_bytes(body[i + 8:i + 16], "big") if ver == 1 else int.from_bytes(body[i + 8:i + 12], "big")
    raise ValueError("no moof")


def listing(files):
    return "".join(f"{hashlib.sha256(b).hexdigest()}  {p}\n" for p, b in sorted(files.items(), key=lambda x: x[0].encode()))


def read_tree(d):
    out = {}
    for root, _, names in os.walk(d):
        for n in names:
            if n.startswith("."):
                continue
            p = os.path.join(root, n)
            out[os.path.relpath(p, d).replace(os.sep, "/")] = open(p, "rb").read()
    return out


def main(argv):
    if len(argv) != 6:
        print(__doc__)
        return 2
    ffmpeg, pc, name, lead, comp = argv[1:]
    # The lead's pieces, from its stage index: each piece's start and the last one's end.
    ix = open(os.path.join(pc, "packs", lead, "package", "index-v.wsix"), "rb").read()
    n = int.from_bytes(ix[16:20], "little")
    starts = [int.from_bytes(ix[160 + 32 * k:168 + 32 * k], "little") for k in range(n)]
    lead_frames = starts[-1] + int.from_bytes(ix[168 + 32 * (n - 1):172 + 32 * (n - 1)], "little")
    total = lead_frames * RATE // FPS
    at = [s * RATE // FPS for s in starts]
    out = os.path.join(pc, "multi", name, "floor", comp)
    if os.path.isdir(out):
        shutil.rmtree(out)
    os.makedirs(out)
    with tempfile.TemporaryDirectory() as tmp:
        src = os.path.join(tmp, "a.wav")
        open(src, "wb").write(wav(total))
        hls = os.path.join(tmp, "hls")
        os.makedirs(hls)
        subprocess.run([ffmpeg, "-nostdin", "-v", "error", "-i", src, "-map", "0:a:0", "-c:a", "libopus", "-b:a", "128k",
                        "-frame_duration", "20", "-application", "audio", "-map_metadata", "-1", "-fflags", "+bitexact",
                        "-flags:a", "+bitexact", "-f", "hls", "-hls_time", f"{STEP}", "-hls_segment_type", "fmp4",
                        "-hls_fmp4_init_filename", "init.mp4", "-hls_segment_filename", os.path.join(hls, "seg_%04d.m4s"),
                        "-hls_playlist_type", "vod", "-hls_list_size", "0", os.path.join(hls, "index.m3u8")], check=True)
        shutil.copyfile(os.path.join(hls, "init.mp4"), os.path.join(out, "init.mp4"))
        frags = [open(os.path.join(hls, f), "rb").read() for f in sorted(os.listdir(hls)) if f.endswith(".m4s")]
        groups = {}
        for f in frags:
            t = tfdt(f)
            k = len(at) if t >= total else max(i for i, s in enumerate(at) if s <= t)
            groups.setdefault(k, []).append(f)
        for k in sorted(groups):
            seg = b""
            for i, f in enumerate(groups[k]):
                for kind, body in boxes(f):
                    if kind == b"styp" and i == 0:
                        seg += body
                    elif kind in (b"moof", b"mdat"):
                        seg += body
            open(os.path.join(out, f"seg_{k:04d}.m4s"), "wb").write(seg)
    print(f"{out}: {sorted(os.listdir(out))}")
    whole = {p: b for p, b in read_tree(os.path.join(pc, "packs", lead, "floor")).items()}
    whole.update({f"{comp}/{p}": b for p, b in read_tree(out).items()})
    digest = hashlib.sha256(listing(whole).encode()).hexdigest()
    p = os.path.join(pc, "multi", name, "policy.toml")
    policy = open(p, encoding="utf-8").read()
    open(p, "w", encoding="utf-8").write(re.sub(r'supplied_listing = "[0-9a-f]{64}"', f'supplied_listing = "{digest}"', policy))
    print(f"listing {digest}")
    return 0


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