# SPDX-License-Identifier: Apache-2.0
"""Synthesise the audio the web-port parity tests run on.

Aimed at the branches in the metric code rather than at sounding like a song:
shorter than each metric's window, digital silence (gated loudness -> -inf /
NaN), a passage below the relative gate, int16 / int24 / float WAV, 96 kHz FLAC
and stereo downmix. When ffmpeg is on PATH it also writes an ID3-tagged MP3 and
Ogg Vorbis / Opus / AAC files for the in-browser decoder checks.

  uv run --no-project --with numpy --with soundfile \
      python tests/parity/fixtures.py OUT_DIR
"""

import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np
import soundfile as sf


def tone_mix(sr: int, seconds: float, seed: int) -> np.ndarray:
  rng = np.random.default_rng(seed)
  t = np.arange(int(sr * seconds)) / sr
  y = (0.3 * np.sin(2 * np.pi * 110 * t)
       + 0.15 * np.sin(2 * np.pi * 1760 * t + 0.3)
       + 0.05 * rng.standard_normal(len(t)))
  # Slow level movement so short-term loudness and LRA have something to show.
  return (y * (0.6 + 0.4 * np.sin(2 * np.pi * 0.07 * t))).astype(np.float32)


def main() -> None:
  out = Path(sys.argv[1] if len(sys.argv) > 1 else "/tmp/ujm-parity")
  out.mkdir(parents=True, exist_ok=True)

  sr = 44100
  left, right = tone_mix(sr, 40, 1), tone_mix(sr, 40, 2) * 0.8
  for ch in (left, right):
    ch[10 * sr:15 * sr] = 0.0                    # below the absolute gate
    ch[20 * sr:25 * sr] *= 10 ** (-45 / 20)      # below the relative gate
  left[30 * sr:30 * sr + 2000] = 0.99            # near-full-scale burst for true peak
  music = out / "music_44k1_s16_stereo.wav"
  sf.write(music, np.stack([left, right], 1), sr, subtype="PCM_16")

  y = tone_mix(48000, 120, 3)
  sf.write(out / "long_48k_s24_stereo.wav", np.stack([y, np.roll(y, 480)], 1), 48000, subtype="PCM_24")
  y = tone_mix(96000, 20, 4)
  sf.write(out / "hires_96k_s24.flac", np.stack([y, -0.5 * y], 1), 96000, subtype="PCM_24")
  sf.write(out / "short_1500ms_48k_float.wav", tone_mix(48000, 1.5, 5), 48000, subtype="FLOAT")
  sf.write(out / "tiny_200ms_44k1_s16.wav", tone_mix(44100, 0.2, 6), 44100, subtype="PCM_16")
  sf.write(out / "silence_5s_48k_s16.wav", np.zeros(5 * 48000, dtype=np.float32), 48000, subtype="PCM_16")

  # Every PCM layout the browser's direct WAV/AIFF reader handles.
  y = tone_mix(48000, 4, 7)
  stereo = np.stack([y, -0.7 * y], 1)
  sf.write(out / "fmt_s24_48k.wav", stereo, 48000, subtype="PCM_24")
  sf.write(out / "fmt_s32_48k.wav", stereo, 48000, subtype="PCM_32")
  sf.write(out / "fmt_f64_48k_mono.wav", y, 48000, subtype="DOUBLE")
  sf.write(out / "fmt_u8_22k05_mono.wav", tone_mix(22050, 4, 8), 22050, subtype="PCM_U8")
  sf.write(out / "fmt_wavex_s24_96k.wav", np.stack([tone_mix(96000, 3, 9)] * 3, 1), 96000,
           format="WAVEX", subtype="PCM_24")
  sf.write(out / "fmt_s16_44k1.aiff", stereo[: 44100 * 4], 44100, format="AIFF", subtype="PCM_16")
  sf.write(out / "fmt_f32_48k.aifc", stereo, 48000, format="AIFF", subtype="FLOAT")

  if shutil.which("ffmpeg"):
    encodes = {
      "tagged_44k1.mp3": ["-codec:a", "libmp3lame", "-b:a", "192k", "-id3v2_version", "3",
                          "-metadata", "artist=ウジ", "-metadata", "title=テスト曲"],
      "vorbis_44k1.ogg": ["-codec:a", "libvorbis", "-q:a", "5"],
      "opus_48k.opus": ["-codec:a", "libopus", "-b:a", "128k"],
      "aac_44k1.m4a": ["-codec:a", "aac", "-b:a", "192k"],
    }
    for name, args in encodes.items():
      try:
        subprocess.run(["ffmpeg", "-loglevel", "error", "-y", "-i", str(music), *args, str(out / name)],
                       check=True)
      except subprocess.CalledProcessError as exc:
        print(f"skipped {name}: ffmpeg failed ({exc.returncode})")

  for f in sorted(out.iterdir()):
    print(f"  {f.name:32s} {f.stat().st_size / 1024:8.0f} KiB")


if __name__ == "__main__":
  main()