#!/usr/bin/env python3
"""Build lhotse CutSets from cleaned manifests (one cut == one utterance).

Reads {train,dev,eval_*}.jsonl (fields: audio,dur,source,tok,text) and writes
cuts_{name}.jsonl.gz. Supervision.text = 'tok' (▁-joined, codepoint units).
All cuts resampled to 16 kHz (on-the-fly). Features are NOT precomputed
(training uses --on-the-fly-feats).
"""
import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor

from lhotse import CutSet, Recording, SupervisionSegment
from lhotse.audio import AudioSource


def make_cut(row, idx):
    from lhotse import MonoCut
    try:
        rec = Recording.from_file(row["audio"], recording_id=f"{row['source']}_{idx:07d}")
    except Exception as e:
        return None
    sup = SupervisionSegment(
        id=rec.id, recording_id=rec.id, start=0.0,
        duration=rec.duration, channel=rec.channel_ids[0],
        text=row["tok"], language=row["source"],
    )
    cut = MonoCut(id=rec.id, start=0.0, duration=rec.duration,
                  channel=rec.channel_ids[0], recording=rec, supervisions=[sup])
    if rec.sampling_rate != 16000:
        cut = cut.resample(16000)
    return cut


def build(jsonl, out, workers):
    rows = [json.loads(l) for l in open(jsonl)]
    cuts = []
    with ThreadPoolExecutor(max_workers=workers) as ex:
        for c in ex.map(lambda a: make_cut(a[1], a[0]), list(enumerate(rows))):
            if c is not None:
                cuts.append(c)
    cs = CutSet.from_cuts(cuts)
    cs.to_file(out)
    dur = sum(c.duration for c in cuts) / 3600
    print(f"{os.path.basename(out)}: {len(cuts)} cuts {dur:.1f}h (from {len(rows)} rows)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifests", required=True, help="dir with *.jsonl")
    ap.add_argument("--out", required=True, help="output dir for cuts")
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--sets", nargs="+",
                    default=["train", "dev", "eval_iv_hi", "eval_svarah", "eval_call_test"])
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    for name in args.sets:
        jsonl = os.path.join(args.manifests, f"{name}.jsonl")
        if not os.path.exists(jsonl):
            print(f"skip {name} (no {jsonl})"); continue
        build(jsonl, os.path.join(args.out, f"cuts_{name}.jsonl.gz"), args.workers)


if __name__ == "__main__":
    main()