"""Phase 2 data prep (H100 node): every source -> MMS forced alignment -> sentence segmentation of long clips ->
punctuation-consistent pause normalisation -> 24 kHz int16 + meta.  Output ~/data2/<ds>/{pcm24k.i16, meta.jsonl, stats.json}

  python -m da.prep2 waxal bible_hau bible_yor slr86 slr70 fleurs
"""
import os, sys, io, json, glob, csv, hashlib, time
import numpy as np, soundfile as sf, librosa
from multiprocessing import Pool
sys.path.insert(0, os.path.expanduser("~/da"))
from da.p2common import words_for_align, Aligner, normalise_pauses   # noqa: E402

RAW = "/scratch/raw"
OUT = os.path.expanduser(os.environ.get("P2OUT", "~/data2"))
OLD = os.path.expanduser("~/data")
SR = 24000
BIBLE_HAU_CAP_H = 30.0
FLOW_OK = {"waxal", "bible_yor", "bible_hau", "slr86", "slr70"}     # 48 kHz studio-ish sources (FLEURS is 16 kHz)


def clean_text(t):
    import unicodedata, re
    t = unicodedata.normalize("NFC", t or "")
    t = t.replace("\u2019", "'").replace("\u2018", "'").replace("\u201c", '"').replace("\u201d", '"')
    return re.sub(r"\s+", " ", t).strip()


def hsplit(i):
    return "val" if int(hashlib.md5(i.encode()).hexdigest(), 16) % 100 == 0 else "train"


# ------------------------------------------------------------------ enumerators (yield meta + loader args)
def items_waxal():
    import pyarrow.parquet as pq
    for lang in ["yor", "ibo", "hau", "pcm"]:
        files = sorted(glob.glob(f"{RAW}/waxal/data/TTS/{lang}/*.parquet"))
        files = [f for f in files if not os.path.basename(f).startswith("yor_tts")]
        for f in files:
            bn = os.path.basename(f)
            split = "test" if "test" in bn else ("val" if "validation" in bn else "train")
            pf = pq.ParquetFile(f)
            for b in pf.iter_batches(batch_size=8):
                for r in b.to_pylist():
                    yield dict(id=f"waxal_{lang}_{r['id']}", ds="waxal", lang=lang,
                               spk=f"waxal_{lang}_{r.get('speaker_id') or 'spk'}", gender=r.get("gender") or "",
                               split=split, text=r["text"]), ("bytes", r["audio"]["bytes"])


FLEURS = {"hau": "ha_ng", "ibo": "ig_ng", "yor": "yo_ng", "eng": "en_us"}


def items_fleurs():
    """raw FLEURS train audio (16 kHz); speaker ids = Phase 1 ECAPA pseudo-speakers from ~/data/fleurs/meta.jsonl"""
    for l in open(f"{OLD}/fleurs/meta.jsonl", encoding="utf-8"):
        m = json.loads(l)
        fn = m["id"].split("_", 2)[2] + ".wav"
        p = f"{RAW}/fleurs/data/{FLEURS[m['lang']]}/audio/train/{fn}"
        yield dict(id=m["id"], ds="fleurs", lang=m["lang"], spk=m["spk"], gender=m.get("gender", ""),
                   split=hsplit(m["id"]), text=m["text"]), ("path", p)


def items_slr(num, lang, prefix):
    for idx in sorted(glob.glob(f"{RAW}/slr{num}/**/line_index.tsv", recursive=True)):
        d = os.path.dirname(idx)
        for line in open(idx, encoding="utf-8"):
            parts = line.rstrip("\n").split("\t")
            if len(parts) < 2:
                continue
            fid, text = parts[0].strip(), parts[-1]
            p = os.path.join(d, fid + ".wav")
            if os.path.exists(p):
                spk = "_".join(fid.split("_")[:2])
                yield dict(id=f"{prefix}_{fid}", ds=prefix, lang=lang, spk=f"{prefix}_{spk}",
                           gender="female" if fid[2] == "f" else "male", split=hsplit(fid), text=text), ("path", p)


def items_bible(lang_dir, lang, ds, cap_h=None):
    root = f"{RAW}/bible/{lang_dir}"
    auds = sorted(glob.glob(f"{root}/**/*.flac", recursive=True) + glob.glob(f"{root}/**/*.wav", recursive=True))
    if cap_h:   # deterministic random subset of the train split, roughly cap_h hours (verse ~ 9 s on average)
        auds.sort(key=lambda p: hashlib.md5(p.encode()).hexdigest())
    tot = 0.0
    for p in auds:
        txt = os.path.splitext(p)[0] + ".txt"
        if not os.path.exists(txt):
            continue
        split = "test" if "/test/" in p else ("val" if "/dev/" in p else "train")
        if cap_h and split == "train":
            if tot > cap_h * 3600:
                continue
            try:
                tot += sf.info(p).duration
            except Exception:
                continue
        yield dict(id=f"{ds}_{os.path.splitext(os.path.basename(p))[0]}", ds=ds, lang=lang, spk=ds, gender="",
                   split=split, text=open(txt, encoding="utf-8").read()), ("path", p)


def enum(ds):
    return {"waxal": items_waxal, "fleurs": items_fleurs,
            "slr86": lambda: items_slr(86, "yor", "slr86"), "slr70": lambda: items_slr(70, "eng", "slr70"),
            "bible_yor": lambda: items_bible("yoruba", "yor", "bible_yor"),
            "bible_hau": lambda: items_bible("hausa", "hau", "bible_hau", BIBLE_HAU_CAP_H)}[ds]()


# ------------------------------------------------------------------ workers
def load(arg):
    meta, src = arg
    try:
        if src[0] == "bytes":
            y, sr = sf.read(io.BytesIO(src[1]), dtype="float32", always_2d=True); y = y.mean(1)
        elif src[0] == "path":
            y, sr = sf.read(src[1], dtype="float32", always_2d=True); y = y.mean(1)
        else:
            pcm = np.memmap(src[1], dtype=np.int16, mode="r")
            y, sr = pcm[src[2]:src[2] + src[3]].astype(np.float32) / 32768.0, 16000
        meta["text"] = clean_text(meta["text"])
        meta["sr_in"] = int(sr)
        if not meta["text"] or len(y) < 0.5 * sr:
            return meta, None, None, "empty"
        y24 = librosa.resample(y, orig_sr=sr, target_sr=SR, res_type="soxr_hq") if sr != SR else y
        y16 = librosa.resample(y, orig_sr=sr, target_sr=16000, res_type="soxr_hq") if sr != 16000 else y
        return meta, y16, y24, None
    except Exception as e:  # noqa
        return meta, None, None, f"err:{type(e).__name__}"


def plan_segments(words, ali, dur):
    """word-index ranges [(a, b)] (inclusive).  Short clips: one segment.  Long clips: whole sentences grouped
    into 3-14 s segments; an over-long sentence is cut at clause/comma marks."""
    n = len(words)
    if dur <= 20.0:
        return [(0, n - 1)]
    sents, a = [], 0
    for i, w in enumerate(words):
        if w["punct"] == "sent" or i == n - 1:
            sents.append((a, i)); a = i + 1
    pieces = []
    for a, b in sents:
        if ali[b]["end"] - ali[a]["start"] <= 16.0:
            pieces.append((a, b)); continue
        s = a
        for i in range(a, b + 1):
            if (words[i]["punct"] in ("clause", "comma") and ali[i]["end"] - ali[s]["start"] > 3.0) or i == b:
                pieces.append((s, i)); s = i + 1
    segs, cur = [], None
    for a, b in pieces:
        if cur is None:
            cur = [a, b]; continue
        if ali[b]["end"] - ali[cur[0]]["start"] <= 14.0 and ali[cur[1]]["end"] - ali[cur[0]]["start"] < 4.0:
            cur[1] = b
        else:
            segs.append(tuple(cur)); cur = [a, b]
    if cur:
        segs.append(tuple(cur))
    return segs


def post(arg):
    meta, y24, words, ali = arg
    out, why = [], {}
    segs = plan_segments(words, ali, len(y24) / SR)
    from da.p2common import silence_mask
    sil = silence_mask(y24, SR)
    for k, (a, b) in enumerate(segs):
        sc = float(np.mean([ali[i]["score"] for i in range(a, b + 1)]))
        maxw = max(ali[i]["end"] - ali[i]["start"] for i in range(a, b + 1))
        text = " ".join(words[i]["word"] for i in range(a, b + 1))
        if sc < 0.30:
            why["low_score"] = why.get("low_score", 0) + 1; continue
        if maxw > 2.5:
            why["long_word"] = why.get("long_word", 0) + 1; continue
        lo = ali[a - 1]["end"] + 0.02 if a > 0 else 0.0
        hi = ali[b + 1]["start"] - 0.02 if b + 1 < len(ali) else None
        y, pauses = normalise_pauses(y24, SR, words, ali, a, b, sil=sil, lo_s=lo, hi_s=hi)
        dur = len(y) / SR
        cps = len(text) / dur
        if not (0.8 <= dur <= 20.0):
            why["dur"] = why.get("dur", 0) + 1; continue
        if not (4.0 <= cps <= 25.0):
            why["cps"] = why.get("cps", 0) + 1; continue
        rms = np.sqrt((y ** 2).mean() + 1e-12)
        y = y * (0.07 / rms)
        pk = np.abs(y).max()
        if pk > 0.99:
            y = y * (0.99 / pk)
        m = dict(meta)
        if len(segs) > 1:
            m["id"] = f"{meta['id']}_s{k:03d}"; m["src_id"] = meta["id"]; m["longform"] = True
        m.update(text=text, dur=round(dur, 3), ascore=round(sc, 3), cps=round(cps, 2),
                 pauses=[round(p, 2) for p in pauses],
                 puncts=[words[i]["punct"] for i in range(a, b)], flow_ok=meta["ds"] in FLOW_OK)
        out.append((m, np.clip(y * 32767, -32768, 32767).astype(np.int16)))
    return out, why


def run(ds):
    od = f"{OUT}/{ds}"
    os.makedirs(od, exist_ok=True)
    al = Aligner()
    stats = dict(src_items=0, src_hours=0.0, kept=0, hours=0.0, rej={}, longform_src=0, longform_hours_in=0.0,
                 longform_segments=0)
    t0 = time.time(); off = 0
    pending = []
    with open(f"{od}/pcm24k.i16", "wb") as pf, open(f"{od}/meta.jsonl", "w", encoding="utf-8") as jf, \
            Pool(20) as lp, Pool(6) as pp:
        def drain(block=False):
            nonlocal off
            keep = []
            for r in pending:
                if block or r.ready():
                    segs, why = r.get()
                    for k, v in why.items():
                        stats["rej"][k] = stats["rej"].get(k, 0) + v
                    for m, pcm in segs:
                        m["off24"] = off; m["n24"] = int(len(pcm))
                        pf.write(pcm.tobytes()); off += len(pcm)
                        jf.write(json.dumps(m, ensure_ascii=False) + "\n")
                        stats["kept"] += 1; stats["hours"] += m["dur"] / 3600
                        if m.get("longform"):
                            stats["longform_segments"] += 1
                else:
                    keep.append(r)
            pending[:] = keep
        import itertools
        src = enum(ds)
        if os.environ.get("LIMIT"):
            src = itertools.islice(src, 0, None, int(os.environ.get("STRIDE", "1")))
            src = itertools.islice(src, int(os.environ["LIMIT"]))
        for meta, y16, y24, err in lp.imap(load, src, chunksize=2):
            stats["src_items"] += 1
            if err:
                stats["rej"][err] = stats["rej"].get(err, 0) + 1; continue
            d = len(y24) / SR
            stats["src_hours"] += d / 3600
            if d > 20:
                stats["longform_src"] += 1; stats["longform_hours_in"] += d / 3600
            if d > 2400:
                stats["rej"]["too_long"] = stats["rej"].get("too_long", 0) + 1; continue
            words = words_for_align(meta["text"])
            ali = al.align(y16, words) if words else None
            if ali is None:
                stats["rej"]["align_fail"] = stats["rej"].get("align_fail", 0) + 1; continue
            pending.append(pp.apply_async(post, ((meta, y24, words, ali),)))
            if len(pending) > 64:
                drain(block=False)
                while len(pending) > 128:
                    time.sleep(0.05); drain(block=False)
            if stats["src_items"] % 2000 == 0:
                print(ds, stats["src_items"], "kept", stats["kept"], "h", round(stats["hours"], 2),
                      "rej", stats["rej"], round(time.time() - t0), "s", flush=True)
        drain(block=True)
    stats = {k: (round(v, 3) if isinstance(v, float) else v) for k, v in stats.items()}
    json.dump(stats, open(f"{od}/stats.json", "w"), indent=1)
    print(ds, "PREP2_DONE", json.dumps(stats), round(time.time() - t0), "s", flush=True)


if __name__ == "__main__":
    for ds in sys.argv[1:]:
        run(ds)