#!/usr/bin/env python3
# Copyright (c) 2026, CIIL Anuvani-LID. Apache-2.0.
"""Deterministic long recordings built from the gate clips, for checking the long-audio pipeline anywhere
the serving kit's gate audio is (the 10-minute and 1-hour k6 files are not distributed).

Every file is a sequence of real clips (suite1k, known languages only) separated by digital silence, laid
out to exercise each rule of the pipeline:
  single_*      one language; gaps of 0.1-3 s (around the VAD's 0.3 s minimum silence and the 0.6 s merge gap)
  mixed_*       three languages in blocks
  continuous_*  clips back to back for over a minute: the VAD's 30 s maximum-speech split and piece tiling
  short_speech  a 0.2 s clip and silence: speech shorter than the 0.3 s minimum span (no decision)
  silence       10 s of zeros (no speech)
  resampled_*   a mixed file at 8 kHz and at 44.1 kHz (resampling; reported, not gated: the server
                resamples with soxr HQ, phones with sherpa-onnx's windowed-sinc resampler)
Each WAV (16-bit PCM) gets a JSON with its ground-truth timeline.

    python3 gate/make_long_audio.py [--out results/tmp/long]
"""
import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tools"))
import common  # noqa: E402

SR = common.SR


def suite_by_language():
    g = common.gate_dir()
    labels = set(common.labels())
    by = {}
    for r in (json.loads(l) for l in open(f"{g}/suite1k/manifest.jsonl")):
        if r["label"] in labels and r["seconds"] >= 1.0:
            by.setdefault(r["label"], []).append(f"{g}/suite1k/{r['file']}")
    return {k: sorted(v) for k, v in sorted(by.items())}


def build(rng, items, gaps):
    """items: [(path, label)], gaps: seconds after each item -> (audio, timeline)"""
    parts, timeline, t = [np.zeros(int(0.5 * SR), np.float32)], [], int(0.5 * SR)
    for (path, label), gap in zip(items, gaps):
        w = common.load_wav(path)
        parts.append(w); timeline.append([t / SR, (t + len(w)) / SR, label]); t += len(w)
        z = np.zeros(int(gap * SR), np.float32)
        parts.append(z); t += len(z)
    return np.concatenate(parts), timeline


def write(out, name, audio, timeline, sr=SR):
    import soundfile as sf
    if sr != SR:
        from fractions import Fraction
        from scipy.signal import resample_poly
        f = Fraction(sr, SR)
        audio = resample_poly(audio, f.numerator, f.denominator).astype(np.float32)
    sf.write(f"{out}/{name}.wav", np.clip(audio, -1, 1), sr, subtype="PCM_16")
    share = {}
    for s, e, l in timeline:
        share[l] = share.get(l, 0.0) + e - s
    meta = {"file": f"{name}.wav", "sample_rate": sr, "seconds": round(len(audio) / sr, 3),
            "dominant": max(share, key=share.get) if share else None, "timeline": timeline}
    json.dump(meta, open(f"{out}/{name}.json", "w"), indent=1)
    print(f"  {name}.wav  {meta['seconds']:7.1f} s  {len(timeline):3d} clips  dominant={meta['dominant']}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(common.REPO, "results", "tmp", "long"))
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    rng = np.random.default_rng(20261008)
    by = suite_by_language()
    langs = list(by)
    gap_choices = [0.1, 0.25, 0.4, 0.55, 0.65, 0.8, 1.5, 3.0]

    def pick(lang, k):
        files = by[lang]
        idx = rng.choice(len(files), size=min(k, len(files)), replace=False)
        return [(files[i], lang) for i in sorted(idx)]

    for i, lang in enumerate(rng.choice(langs, size=6, replace=False)):
        items = pick(lang, 24)
        audio, tl = build(rng, items, rng.choice(gap_choices, size=len(items)))
        write(a.out, f"single_{i}_{lang}", audio, tl)

    mixed = []
    for i in range(4):
        items = []
        for lang in rng.choice(langs, size=3, replace=False):
            items += pick(lang, 10)
        audio, tl = build(rng, items, rng.choice(gap_choices, size=len(items)))
        write(a.out, f"mixed_{i}", audio, tl)
        mixed.append((audio, tl))

    for i in range(2):
        items = []
        for lang in rng.choice(langs, size=2, replace=False):
            items += pick(lang, 14)
        audio, tl = build(rng, items, [0.0] * len(items))
        write(a.out, f"continuous_{i}", audio, tl)

    g = common.gate_dir()
    short = [r for r in (json.loads(l) for l in open(f"{g}/suite1k/manifest.jsonl")) if r["seconds"] <= 0.25][0]
    w = common.load_wav(f"{g}/suite1k/{short['file']}")
    audio = np.concatenate([np.zeros(2 * SR, np.float32), w, np.zeros(3 * SR, np.float32)])
    write(a.out, "short_speech", audio, [[2.0, 2.0 + len(w) / SR, short["label"]]])
    write(a.out, "silence", np.zeros(10 * SR, np.float32), [])

    for sr in (8000, 44100):
        write(a.out, f"resampled_{sr}", mixed[0][0], mixed[0][1], sr=sr)


if __name__ == "__main__":
    main()