#!/usr/bin/env python3
"""Write the wt3d reference files next to this script.

Independent of WaveStack's encoder on purpose: these files check that an
implementation (WaveStack's included) reads the format as the spec says.
Every frame is a constant level equal to its index / 100, so a reader can
check which frame landed where. Run: python3 wt3d_vectors.py

Specification: https://wavestack.app/wt3d/ (public domain, CC0).
"""

import struct
from pathlib import Path

OUT = Path(__file__).resolve().parent


def chunk(cid: bytes, body: bytes) -> bytes:
    pad = b"\0" if len(body) % 2 else b""
    return cid + struct.pack("<I", len(body)) + body + pad


def wt3d(frame_size, cols, rows, layers, flags=0, names=(), version=1, extra=b""):
    block = b"".join(n.encode("utf-8") + b"\0" for n in names)
    head = struct.pack("<7I", version, frame_size, cols, rows, layers, flags, len(block))
    return chunk(b"wt3d", head + block + extra)


def wav(frame_size, frames, *chunks):
    """Mono 32-bit float WAV at 44.1 kHz with a clm tag and the given chunks."""
    samples = b"".join(struct.pack("<f", i / 100) * frame_size for i in range(frames))
    fmt = chunk(b"fmt ", struct.pack("<HHIIHH", 3, 1, 44100, 44100 * 4, 4, 32))
    clm = chunk(b"clm ", f"<!>{frame_size:04} 10000000 wavetable (wt3d reference)".encode())
    body = b"WAVE" + fmt + clm + b"".join(chunks) + chunk(b"data", samples)
    return b"RIFF" + struct.pack("<I", len(body)) + body


VECTORS = {
    # 4 columns x 2 rows, one layer.
    "grid-4x2.wav": wav(64, 8, wt3d(64, 4, 2, 1)),
    # 3 x 2 x 2 volume with layer names and hint flags (normalized, zero-aligned).
    "volume-3x2x2.wav": wav(64, 12, wt3d(64, 3, 2, 2, flags=0b101, names=("Glass", "Reed"))),
    # A later version that appends fields: read as 2 x 2.
    "future-version.wav": wav(64, 4, wt3d(64, 2, 2, 1, version=2, extra=b"\x09" * 8)),
    # Claims 4 x 3 = 12 frames but holds 8: must be ignored (opens as one row of 8).
    "invalid-count.wav": wav(64, 8, wt3d(64, 4, 3, 1)),
}

if __name__ == "__main__":
    OUT.mkdir(parents=True, exist_ok=True)
    for name, data in VECTORS.items():
        (OUT / name).write_bytes(data)
        print(f"{name}: {len(data)} bytes")