#!/usr/bin/env python3
"""Convert the 4 Egyptian HF datasets into fish-speech finetune layout.

Output layout (symlinked wavs to save disk; fap loudness-norm copies later):
  /opt/work/data-raw-fs/<speaker>/<clip>.wav  (symlink)
  /opt/work/data-raw-fs/<speaker>/<clip>.lab  (transcript)

Speaker grouping:
  - noselleel (multi-voice): per-video pseudo speaker  nos_<VIDEOID>
  - eqkawkab  (single narrator, ECAPA-isolated):        eqk
  - moustafa-sadek:                                     sad_<VIDEOID>
  - mosaifside:                                         mos_<VIDEOID>

Text policy (voice agent will emit PLAIN text, but tashkeel disambiguates):
  - datasets with diacritized text: deterministic 50/50 per clip -> half keep
    tashkeel, half stripped (model learns both input styles)
  - plain datasets: as-is
Dev/holdout entries are excluded from training and dumped to /opt/work/eval/.
"""

import argparse
import hashlib
import json
import os
import re
import sys
import unicodedata
from pathlib import Path

RAW = Path("/opt/work/data_raw")
EVAL = Path("/opt/work/eval")
# set in main() from CLI args (per-stage roots for incremental processing)
OUT = None
OUT_VAL = None

# harakat, Quranic annotation marks, superscript alef, tatweel
TASHKEEL_RE = re.compile(r"[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭـ]")
WS_RE = re.compile(r"\s+")


def strip_tashkeel(t: str) -> str:
    return WS_RE.sub(" ", TASHKEEL_RE.sub("", t)).strip()


def clean(t: str) -> str:
    t = unicodedata.normalize("NFC", t)
    t = "".join(c for c in t if unicodedata.category(c)[0] != "C")
    return WS_RE.sub(" ", t).strip()


def keep_diac(clip_id: str) -> bool:
    return int(hashlib.md5(clip_id.encode()).hexdigest(), 16) % 2 == 0


def video_id(clip_id: str) -> str:
    # clip ids look like  -YVpVr8AF18_0007  ->  video -YVpVr8AF18
    return clip_id.rsplit("_", 1)[0]


def link(wav: Path, spk: str, clip_id: str, text: str, stats: dict, root: Path = None):
    d = (root or OUT) / spk
    d.mkdir(parents=True, exist_ok=True)
    lnk = d / f"{clip_id}.wav"
    if not lnk.exists():
        lnk.symlink_to(wav.resolve())
    (d / f"{clip_id}.lab").write_text(text, encoding="utf-8")
    stats["clips"] += 1
    stats["chars"] += len(text)


def load_jsonl(p: Path):
    with open(p, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                yield json.loads(line)


def audio_path(ds_dir: Path, rec: dict) -> Path | None:
    for k in ("audio_path", "audio", "path", "file"):
        if k in rec:
            p = ds_dir / rec[k]
            if p.exists():
                return p
            p2 = ds_dir / "audio" / Path(rec[k]).name
            if p2.exists():
                return p2
    if "id" in rec:
        p = ds_dir / "audio" / f"{rec['id']}.wav"
        if p.exists():
            return p
    return None


def process(name, spk_fn, diac_mix, train_files, dev_files, text_key="text"):
    ds = RAW / name
    stats = {"clips": 0, "chars": 0, "missing": 0, "empty": 0}
    dev_rows = []
    for fname, is_dev in [(f, False) for f in train_files] + [(f, True) for f in dev_files]:
        p = ds / fname
        if not p.exists():
            print(f"  !! {name}: manifest {fname} missing", file=sys.stderr)
            continue
        for rec in load_jsonl(p):
            cid = rec.get("id") or Path(rec.get("audio", rec.get("audio_path", ""))).stem
            text = clean(str(rec.get(text_key) or ""))
            if not text or len(text) < 2:
                stats["empty"] += 1
                continue
            wav = audio_path(ds, rec)
            if wav is None:
                stats["missing"] += 1
                continue
            if is_dev:
                dev_rows.append({"id": cid, "text": text, "wav": str(wav), "ds": name})
                link(wav, spk_fn(cid), cid, text, stats, root=OUT_VAL)
                stats["clips"] -= 1  # don't count val clips as train
                continue
            if diac_mix and not keep_diac(cid):
                text = strip_tashkeel(text)
            link(wav, spk_fn(cid), cid, text, stats)
    print(f"{name}: {stats}")
    return dev_rows


DATASETS = {
    "noselleel": ("noselleel-egyptian-tts", lambda c: f"nos_{video_id(c)}", True,
                  ["train.jsonl"], ["dev.jsonl"]),
    "eqkawkab": ("eqkawkab-egyptian-tts", lambda c: "eqk", True,
                 ["train.jsonl"], ["dev.jsonl"]),
    "sadek": ("moustafa-sadek-egyptian-tts", lambda c: f"sad_{video_id(c)}", False,
              ["transcripts.jsonl"], []),
    "mosaif": ("mosaifside-egyptian-tts", lambda c: f"mos_{video_id(c)}", False,
               ["transcripts.jsonl"], []),
}


def main():
    global OUT, OUT_VAL
    ap = argparse.ArgumentParser()
    ap.add_argument("--datasets", required=True,
                    help="comma list from: " + ",".join(DATASETS))
    ap.add_argument("--out", required=True)
    ap.add_argument("--out-val", required=True)
    ap.add_argument("--holdout-suffix", default="")
    args = ap.parse_args()
    OUT, OUT_VAL = Path(args.out), Path(args.out_val)
    OUT.mkdir(parents=True, exist_ok=True)
    OUT_VAL.mkdir(parents=True, exist_ok=True)
    EVAL.mkdir(parents=True, exist_ok=True)
    dev = []
    for key in args.datasets.split(","):
        dev += process(*DATASETS[key.strip()])

    holdout = EVAL / f"dev_holdout{args.holdout_suffix}.jsonl"
    with open(holdout, "w", encoding="utf-8") as f:
        for r in dev:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    spk = len([d for d in OUT.iterdir() if d.is_dir()])
    print(f"TOTAL speakers={spk} dev_holdout={len(dev)}")


if __name__ == "__main__":
    main()