#!/usr/bin/env python3
# Published copy, nothing is changed except this line. The sha256 of the original is in sha256.txt.
"""Build the ~1 h test file from LibriSpeech test-clean (CC BY 4.0).

Deterministic: whole chapters, sorted by (speaker id, chapter id) as integers, greedily taken while
the running total stays <= LIMIT_S. Utterances of one chapter are joined in corpus order with
GAP_UTT_S of digital silence, chapters with GAP_CHAP_S. Output: 16 kHz mono PCM16 WAV.
Also writes the reference (raw LibriSpeech text, upper case, no punctuation), a manifest with the
start/end second of every utterance, and a chunked variant (fallback for 16 GB: chunks <= CHUNK_S,
cut only at utterance boundaries).

usage: build_audio.py <LibriSpeech/test-clean dir> <out dir>
"""
import hashlib
import json
import sys
from pathlib import Path

import numpy as np
import soundfile as sf

SR = 16000
LIMIT_S = 3590.0      # stay below the model's one-hour limit (policy duration limit 3600 s)
GAP_UTT_S = 0.25
GAP_CHAP_S = 1.0
CHUNK_S = 400.0       # 8192-token context: 400 s * 12.5 + 400 s * 6 + 128 = 7528 < 8192


def sha256(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    return h.hexdigest()


def main(src: Path, out: Path) -> None:
    out.mkdir(parents=True, exist_ok=True)
    chapters = []
    for spk in sorted((p for p in src.iterdir() if p.is_dir()), key=lambda p: int(p.name)):
        for chap in sorted((p for p in spk.iterdir() if p.is_dir()), key=lambda p: int(p.name)):
            trans = chap / f"{spk.name}-{chap.name}.trans.txt"
            utts = []
            for line in trans.read_text(encoding="utf-8").splitlines():
                uid, text = line.split(" ", 1)
                utts.append((uid, text.strip()))
            dur = 0.0
            for uid, _ in utts:
                info = sf.info(str(chap / f"{uid}.flac"))
                assert info.samplerate == SR and info.channels == 1, uid
                dur += info.frames / SR
            dur += GAP_UTT_S * (len(utts) - 1)
            chapters.append({"speaker": spk.name, "chapter": chap.name, "dir": chap, "utts": utts, "dur": dur})

    chosen, total = [], 0.0
    for c in chapters:
        add = c["dur"] + (GAP_CHAP_S if chosen else 0.0)
        if total + add <= LIMIT_S:
            chosen.append(c)
            total += add
    print(f"chapters available {len(chapters)}, chosen {len(chosen)}, total {total:.2f} s")

    pieces, manifest_utts, ref_lines, pos = [], [], [], 0
    chap_manifest = []
    for ci, c in enumerate(chosen):
        if ci:
            pieces.append(np.zeros(int(GAP_CHAP_S * SR), dtype=np.int16))
            pos += int(GAP_CHAP_S * SR)
        c_start = pos
        for ui, (uid, text) in enumerate(c["utts"]):
            if ui:
                pieces.append(np.zeros(int(GAP_UTT_S * SR), dtype=np.int16))
                pos += int(GAP_UTT_S * SR)
            x, sr = sf.read(str(c["dir"] / f"{uid}.flac"), dtype="int16")
            assert sr == SR
            manifest_utts.append({"id": uid, "chapter": f'{c["speaker"]}-{c["chapter"]}', "start_s": round(pos / SR, 3),
                                  "end_s": round((pos + len(x)) / SR, 3), "text": text})
            ref_lines.append(text)
            pieces.append(x)
            pos += len(x)
        chap_manifest.append({"chapter": f'{c["speaker"]}-{c["chapter"]}', "speaker": c["speaker"], "start_s": round(c_start / SR, 3),
                              "end_s": round(pos / SR, 3), "n_utts": len(c["utts"]),
                              "n_words": sum(len(t.split()) for _, t in c["utts"])})
    audio = np.concatenate(pieces)
    dur = len(audio) / SR
    wav = out / "librispeech-test-clean-1h.wav"
    sf.write(str(wav), audio, SR, subtype="PCM_16")
    (out / "reference-librispeech.txt").write_text("\n".join(ref_lines) + "\n", encoding="utf-8")
    n_words = sum(len(t.split()) for t in ref_lines)

    # chunked variant: cut only between utterances, chunk <= CHUNK_S
    chunks_dir = out / "chunks"
    chunks_dir.mkdir(exist_ok=True)
    chunks, cur, cur_start = [], [], None
    for u in manifest_utts:
        if cur and (u["end_s"] - cur_start) > CHUNK_S:
            chunks.append(cur)
            cur, cur_start = [], None
        if cur_start is None:
            cur_start = u["start_s"]
        cur.append(u)
    if cur:
        chunks.append(cur)
    chunk_manifest = []
    for i, ch in enumerate(chunks):
        a, b = int(ch[0]["start_s"] * SR), int(ch[-1]["end_s"] * SR)
        p = chunks_dir / f"chunk-{i:02d}.wav"
        sf.write(str(p), audio[a:b], SR, subtype="PCM_16")
        (chunks_dir / f"chunk-{i:02d}.ref.txt").write_text("\n".join(u["text"] for u in ch) + "\n", encoding="utf-8")
        chunk_manifest.append({"file": p.name, "start_s": ch[0]["start_s"], "end_s": ch[-1]["end_s"],
                               "dur_s": round((b - a) / SR, 3), "n_utts": len(ch),
                               "n_words": sum(len(u["text"].split()) for u in ch)})

    manifest = {
        "source": "LibriSpeech test-clean (Panayotov et al. 2015), https://www.openslr.org/12/, CC BY 4.0, tarball md5 32fa31d27d2e1cad72775fee3f4849a9",
        "sample_rate": SR, "channels": 1, "format": "PCM16 WAV",
        "gap_between_utterances_s": GAP_UTT_S, "gap_between_chapters_s": GAP_CHAP_S,
        "limit_s": LIMIT_S, "duration_s": round(dur, 3), "n_chapters": len(chosen),
        "n_utterances": len(manifest_utts), "n_ref_words": n_words,
        "wav_sha256": sha256(wav), "reference_sha256": sha256(out / "reference-librispeech.txt"),
        "chapters": chap_manifest, "utterances": manifest_utts, "chunks": chunk_manifest,
    }
    (out / "manifest.json").write_text(json.dumps(manifest, indent=1), encoding="utf-8")
    print(f"duration {dur:.2f} s ({dur/60:.2f} min), utterances {len(manifest_utts)}, ref words {n_words}, chunks {len(chunks)}")
    print("max chunk", max(c["dur_s"] for c in chunk_manifest), "s")
    print("wav sha256", manifest["wav_sha256"])


if __name__ == "__main__":
    main(Path(sys.argv[1]), Path(sys.argv[2]))