"""
Step 2: label the downloaded recordings and split them into train / valid / test manifests.

    python prepare.py --config config/pilot.yaml

Writes data/manifests/<name>/{train,valid,test}.jsonl and tokens.txt, and prints the hours per
reciter. The first run decodes every recording once to measure its true duration and listens to
each ayah 1 for a basmala; both are cached (data/durations.json, data/formulas.json), so re-running
it after changing the data section of the config is quick. Stopping it keeps the work done so far.
"""
import json
import os
import random
import shutil
import sys
import time
import urllib.request
import zlib
from multiprocessing import Pool

from quran_asr import audio, config as config_mod
from quran_asr import quran, text

NVIDIA_URL = ("https://huggingface.co/mah92/sherpa-onnx-nemo-ctc-ar-fastconformer-hybrid-large-pcd-v1.0-"
              "non-streaming-int8-2024-12-25/resolve/main/")


def progress(k, total, started, every):
    """Prints 'k/total, rate, time left' every `every` items and at the end."""
    if k % every and k != total:
        return
    rate = k / max(1e-6, time.time() - started)
    left = (total - k) / max(rate, 1e-6)
    print(f"  {k}/{total}  {rate:.0f} files/s  "
          f"{f'~{left / 60:.0f} min left' if left >= 90 else f'~{left:.0f} s left' if k < total else 'done'}",
          flush=True)


def save_json(path, value):
    with open(path + ".tmp", "w", encoding="utf-8") as f:
        json.dump(value, f)
    os.replace(path + ".tmp", path)


def duration_of(path):
    try:
        return path, audio.decoded_seconds(path)
    except Exception:  # can't be decoded to the end (damaged download)
        return path, None


def load_durations(cfg, paths):
    """Each recording's true duration (None if damaged), measured by decoding it once and cached."""
    cache_path = os.path.join(cfg.paths.data_dir, "durations.json")
    saved = json.load(open(cache_path, encoding="utf-8")) if os.path.exists(cache_path) else {}
    if saved.get("method") != "decoded":
        if saved:
            print("measuring durations by decoding: MP3 headers are wrong for some reciters (once, ~10 min)")
        saved = {"method": "decoded", "durations": {}}
    cache = saved["durations"]
    todo = [p for p in paths if p not in cache]
    print(f"durations: {len(paths) - len(todo)} cached, {len(todo)} to measure")
    if todo:
        started = time.time()
        try:
            with Pool(cfg.data.prepare_workers) as pool:
                for k, (p, d) in enumerate(pool.imap_unordered(duration_of, todo, chunksize=16), 1):
                    cache[p] = d
                    progress(k, len(todo), started, 5000)
                    if k % 50000 == 0:
                        save_json(cache_path, saved)
        finally:
            save_json(cache_path, saved)  # also when interrupted: the next run continues from here
    return cache


# --- Basmala / isti'adha detection ---------------------------------------------------------

def _nvidia_model(cfg):
    """The app's current model (NVIDIA FastConformer Arabic), used only to listen for formulas."""
    folder = os.path.join(cfg.paths.data_dir, "nvidia_fastconformer")
    os.makedirs(folder, exist_ok=True)
    model, tokens = os.path.join(folder, "model.onnx"), os.path.join(folder, "tokens.txt")
    if not os.path.exists(tokens):
        urllib.request.urlretrieve(NVIDIA_URL + "tokens.txt", tokens)
    if not os.path.exists(model):
        print("downloading the NVIDIA model for basmala detection (174 MB)...")
        part = model + ".part"
        urllib.request.urlretrieve(NVIDIA_URL + "model.int8.onnx", part)
        # The export lacks sherpa's metadata; ONNX is protobuf, so appending the fields adds them.
        def field(tag, payload):
            n, out = len(payload), bytearray([tag])
            while True:
                b, n = n & 0x7F, n >> 7
                out.append(b | 0x80 if n else b)
                if not n:
                    break
            return bytes(out) + payload
        with open(part, "ab") as f:
            for k, v in (("vocab_size", "1025"), ("subsampling_factor", "8"), ("normalize_type", "per_feature")):
                f.write(field(0x72, field(0x0A, k.encode()) + field(0x12, v.encode())))
        shutil.move(part, model)
    return model, tokens


_recognizer = None


def _start_detector(model, tokens):
    """Runs once in each worker process: one recognizer per process, one thread each."""
    global _recognizer
    import sherpa_onnx

    _recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(model=model, tokens=tokens, num_threads=1)


def _detect(job):
    path, label = job
    return path, detect_formulas(_recognizer, path, label)


def detect_all(cfg, jobs, cache, cache_path):
    """Listens to the start of every ayah 1 recording in parallel; results are cached as it goes."""
    model, tokens = _nvidia_model(cfg)
    workers = max(1, cfg.data.prepare_workers)
    print(f"listening for basmala / isti'adha in {len(jobs)} ayah 1 recordings ({workers} processes)...")
    started = time.time()
    try:
        with Pool(workers, _start_detector, (model, tokens)) as pool:
            for k, (path, result) in enumerate(pool.imap_unordered(_detect, jobs, chunksize=2), 1):
                cache[path] = result
                progress(k, len(jobs), started, 250)
                if k % 1000 == 0:
                    save_json(cache_path, cache)
    finally:
        save_json(cache_path, cache)  # also when interrupted: the next run continues from here


def _starts_with(heard, phrase, start=0):
    """How many words of `phrase` the heard words follow from position `start` (fuzzy)."""
    words = phrase.split()
    n = 0
    for k, w in enumerate(words):
        if start + k >= len(heard):
            break
        h = heard[start + k]
        if h == w or text.edit_distance(h, w) <= max(1, len(w) // 3):
            n += 1
    return n


def detect_formulas(recognizer, path, label):
    """Returns (istiadha, basmala): whether they are recited before the ayah in this file."""
    try:
        wave = audio.load(path, max_seconds=14)  # formulas come first: decode only the start
    except Exception:
        return False, False  # unreadable: the caller's duration check normally drops it first
    stream = recognizer.create_stream()
    stream.accept_waveform(audio.SAMPLE_RATE, wave)
    recognizer.decode_stream(stream)
    heard = text.normalize(stream.result.text).split()
    pos, istiadha, basmala = 0, False, False
    if _starts_with(heard, text.normalize(quran.ISTIADHA)) >= 3:
        istiadha, pos = True, 5
    if not label.startswith(quran.BASMALA) and _starts_with(heard, quran.BASMALA, pos) >= 3:
        basmala = True
    return istiadha, basmala


# --- Warsh numbering ------------------------------------------------------------------------

def detect_numbering(cfg, recordings, durations, hafs_mapped, native):
    """
    "hafs" or "warsh": how a Warsh reciter's files are numbered on everyayah. Most use Hafs ayah
    numbers, but some (Abdul Basit's Warsh) use the Warsh count, which differs in many surahs, so
    the same file name holds a different ayah. Decided by listening to 30 recordings whose two
    possible labels differ, and seeing which one the reciter actually recites.
    """
    import sherpa_onnx

    model, tokens = _nvidia_model(cfg)
    recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(model=model, tokens=tokens, num_threads=8)
    informative = [(s, a, p) for s, a, p in recordings if 2 <= (durations.get(p) or 0) <= 25
                   and hafs_mapped.get((s, a)) and native.get((s, a)) and hafs_mapped[(s, a)] != native[(s, a)]]
    votes = {"hafs": 0, "warsh": 0}
    for s, a, path in random.Random(0).sample(informative, min(30, len(informative))):
        stream = recognizer.create_stream()
        stream.accept_waveform(audio.SAMPLE_RATE, audio.load(path))
        recognizer.decode_stream(stream)
        heard = text.normalize(stream.result.text)

        def error(label):
            counts = text.error_counts(heard, label)
            return counts[4] / max(1, counts[5])

        votes["warsh" if error(native[(s, a)]) < error(hafs_mapped[(s, a)]) else "hafs"] += 1
    return ("warsh" if votes["warsh"] > votes["hafs"] else "hafs"), votes


def load_tlog(cfg):
    """
    The tlog clips download.py saved: [{path, shard, sura, aya, text}] (labels not yet checked),
    minus the recordings in the Quran-Lab benchmark (data/quranlab/, see benchmark.py).
    """
    folder = os.path.join(cfg.paths.data_dir, "tlog")
    out = []
    if not os.path.isdir(folder):
        return out
    bench = os.path.join(cfg.paths.data_dir, "quranlab", "benchmark.jsonl")
    held_out = set()
    if os.path.exists(bench):
        held_out = {json.loads(l)["id"].split("_")[-1] for l in open(bench, encoding="utf-8") if '"tlog_holdout"' in l}
    for name in sorted(os.listdir(folder)):
        if name.endswith(".jsonl"):
            with open(os.path.join(folder, name), encoding="utf-8") as f:
                for line in f:
                    e = json.loads(line)
                    if os.path.splitext(e["file"])[0].split("_")[-1] in held_out:
                        continue
                    e["path"] = os.path.join(cfg.paths.data_dir, "audio", "tlog", f"{e['shard']:05d}", e["file"])
                    out.append(e)
    return out


_PREFIXES = ("", quran.BASMALA + " ", text.normalize(quran.ISTIADHA) + " ",
             text.normalize(quran.ISTIADHA) + " " + quran.BASMALA + " ")


def transcribe_tlog(cfg, entries):
    """
    What the current model hears on each tlog clip: {path: text}, cached in data/tlog_heard.json,
    so the label rules below can change (or a later, better model can re-check rejected clips
    after deleting the file) without transcribing again. Runs on the GPU (minutes for 200k
    clips): pause training first.
    """
    import torch
    from torch.utils.data import DataLoader

    import export
    from quran_asr import data, model as model_mod

    cache_path = os.path.join(cfg.paths.data_dir, "tlog_heard.json")
    cache = json.load(open(cache_path, encoding="utf-8")) if os.path.exists(cache_path) else {}
    todo = [e for e in entries if e["path"] not in cache]
    print(f"tlog: {len(entries) - len(todo)} clips already transcribed, {len(todo)} to transcribe")
    if not todo:
        return cache
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    state, model_config = export.load_averaged(cfg.paths.run_dir, cfg.export.average_last)
    model = model_mod.build(config_mod._wrap(model_config), len(text.VOCAB))
    model.load_state_dict(state)
    model.to(device).eval()
    loader = DataLoader(data.RecitationDataset(todo), batch_sampler=data.DurationBatchSampler(todo, 400, 1.0, False),
                        collate_fn=data.collate, num_workers=cfg.data.prepare_workers)
    started, done = time.time(), 0
    try:
        with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            for x, lens, _t, _tl, idx in loader:
                log_probs, out_lens = model(x.to(device), lens.to(device))
                for ids, i in zip(model_mod.greedy_decode(log_probs, out_lens), idx):
                    cache[todo[i]["path"]] = text.decode(ids)
                done += len(idx)
                if done // 5000 != (done - len(idx)) // 5000:
                    progress(done, len(todo), started, 1)
                    save_json(cache_path, cache)
    finally:
        save_json(cache_path, cache)
    progress(len(todo), len(todo), started, 1)
    return cache


def best_label(entry, heard, hafs):
    """
    (skeleton error, label) for what this tlog clip most likely contains: its file's ayah, or a
    neighbouring ayah, or two consecutive ayahs (users mis-number and record across a boundary),
    with a basmala and/or isti'adha said first. Spelling-insensitive (skeleton letters).
    """
    from rapidfuzz.distance import Levenshtein

    def skel(t):
        return text.skeleton(t.replace(" ", ""))

    sura, aya = entry["sura"], entry["aya"]
    texts = [entry["text"]]
    for lo, hi in ((aya - 1, aya - 1), (aya + 1, aya + 1), (aya - 1, aya), (aya, aya + 1)):
        if (sura, lo) in hafs and (sura, hi) in hafs:
            texts.append(" ".join(quran.hafs_label(hafs, sura, k) for k in range(lo, hi + 1)))
    said = skel(heard)
    best = (1.0, entry["text"])
    for t in texts:  # the file's own label first: it wins ties
        for prefix in _PREFIXES:
            ref = skel(prefix + t)
            err = Levenshtein.distance(said, ref) / max(1, len(ref))
            if err < best[0]:
                best = (err, prefix + t)
    return best


def load_iqraeval(cfg):
    """The Quranic Iqra'Eval clips download.py kept: [{path, split, sura, aya, text}]."""
    folder = os.path.join(cfg.paths.data_dir, "iqraeval")
    out = []
    if not os.path.isdir(folder):
        return out
    for name in sorted(os.listdir(folder)):
        if name.endswith(".jsonl"):
            with open(os.path.join(folder, name), encoding="utf-8") as f:
                for line in f:
                    e = json.loads(line)
                    e["path"] = os.path.join(cfg.paths.data_dir, "audio", "iqraeval", e["split"], e["file"])
                    out.append(e)
    return out


# --- Main ------------------------------------------------------------------------------------

def main():
    args = config_mod.from_args(__doc__).parse_args()
    cfg = config_mod.load(args.config)
    d = cfg.data
    hafs = quran.load_hafs(cfg.paths.data_dir)
    counts = quran.ayah_counts(hafs)
    warsh = quran.load_warsh(cfg.paths.warsh_text)
    hafs_mapped = quran.warsh_for_hafs(hafs, warsh)  # Warsh text cut at Hafs ayah boundaries
    native = quran.warsh_by_ayah(warsh)  # Warsh text by the Warsh ayah numbers
    warsh_counts = {}
    for s, a in native:
        warsh_counts[s] = max(warsh_counts.get(s, 0), a)
    surahs = list(range(1, 115)) if d.surahs == "all" else list(d.surahs)

    items = []  # (folder, riwaya, role, sura, aya, path)
    for role, group_riwaya, folders in (("train", "hafs", d.hafs_reciters), ("train", "warsh", d.warsh_reciters),
                                        ("test", None, d.test_reciters)):
        for folder in folders:
            # Test folders under warsh/ (e.g. your own recordings) are Warsh; the rest are Hafs.
            riwaya = group_riwaya or ("warsh" if folder.startswith("warsh/") else "hafs")
            folder_dir = os.path.dirname(config_mod.audio_path(cfg.paths.data_dir, folder, 1, 1))
            # One directory listing per reciter (Windows gives sizes for free) instead of a check per file.
            on_disk = set()
            if os.path.isdir(folder_dir):
                on_disk = {e.name for e in os.scandir(folder_dir) if e.is_file() and e.stat().st_size > 0}
            for sura in surahs:
                last = max(counts[sura], warsh_counts[sura]) if riwaya == "warsh" else counts[sura]
                for aya in range(1, last + 1):
                    if f"{sura:03d}{aya:03d}.mp3" in on_disk:
                        path = config_mod.audio_path(cfg.paths.data_dir, folder, sura, aya)
                        items.append((folder, riwaya, role, sura, aya, path))
    print(f"{len(items)} recordings on disk")
    durations = load_durations(cfg, [it[5] for it in items])

    # Which numbering each Warsh reciter's files follow (detected once and cached; the config's
    # optional data.warsh_numbering, e.g. {warsh/some_reciter: warsh}, overrides it).
    numbering_path = os.path.join(cfg.paths.data_dir, "numbering.json")
    numbering = json.load(open(numbering_path, encoding="utf-8")) if os.path.exists(numbering_path) else {}
    for folder in sorted({it[0] for it in items if it[1] == "warsh"}):
        if folder not in numbering and d.basmala_detection == "auto":
            recordings = [(it[3], it[4], it[5]) for it in items if it[0] == folder]
            numbering[folder], votes = detect_numbering(cfg, recordings, durations, hafs_mapped, native)
            print(f"{folder}: files follow the {numbering[folder]} ayah numbering (votes {votes})")
            save_json(numbering_path, numbering)
    numbering.update(vars(d.warsh_numbering) if getattr(d, "warsh_numbering", None) else {})
    labels_by_folder = {f: native if n == "warsh" else hafs_mapped for f, n in numbering.items()}

    def check(folder, riwaya, sura, aya, path):
        """(label, None) for a usable recording, else ("", reason it is skipped)."""
        dur = durations.get(path)
        if dur is None or dur < 0.1:
            return "", "damaged"  # an empty or broken file on everyayah
        if not (d.min_duration <= dur <= d.max_duration):
            return "", "duration"
        if riwaya == "warsh":
            label = labels_by_folder.get(folder, hafs_mapped).get((sura, aya), "")
        else:
            label = quran.hafs_label(hafs, sura, aya)
        return (label, None) if label else ("", "label")

    detect = d.basmala_detection == "auto"
    formula_cache_path = os.path.join(cfg.paths.data_dir, "formulas.json")
    formula_cache = json.load(open(formula_cache_path, encoding="utf-8")) if os.path.exists(formula_cache_path) else {}
    if detect:
        jobs = []
        for folder, riwaya, _role, sura, aya, path in items:
            if aya == 1 and path not in formula_cache:
                label, _ = check(folder, riwaya, sura, aya, path)
                if label:
                    jobs.append((path, label))
        print(f"basmala / isti'adha: {sum(1 for it in items if it[4] == 1) - len(jobs)} ayah 1 recordings "
              f"cached, {len(jobs)} to listen to")
        if jobs:
            detect_all(cfg, jobs, formula_cache, formula_cache_path)

    entries, skipped = [], {"damaged": 0, "duration": 0, "label": 0, "too_fast": 0}
    for folder, riwaya, role, sura, aya, path in items:
        label, reason = check(folder, riwaya, sura, aya, path)
        if reason:
            skipped[reason] += 1
            continue
        dur = durations[path]
        if aya == 1 and detect:
            istiadha, basmala = formula_cache[path]
            if basmala:
                label = quran.BASMALA + " " + label
            if istiadha:
                label = text.normalize(quran.ISTIADHA) + " " + label
        # CTC needs at least one output frame (40 ms) per letter, with room to spare.
        if len(label) * 1.3 > dur * 25:
            skipped["too_fast"] += 1
            continue
        entries.append(dict(path=path, duration=round(dur, 3), text=label, reciter=folder, riwaya=riwaya,
                            role=role, sura=sura, aya=aya))

    rnd = random.Random(cfg.train.seed)
    # Cap hours per training reciter with a random selection of ayahs.
    if d.max_hours_per_reciter:
        capped = []
        for folder in d.hafs_reciters + d.warsh_reciters:
            mine = [e for e in entries if e["reciter"] == folder and e["role"] == "train"]
            rnd.shuffle(mine)
            total = 0.0
            for e in mine:
                if total / 3600 >= d.max_hours_per_reciter:
                    break
                capped.append(e)
                total += e["duration"]
        entries = capped + [e for e in entries if e["role"] == "test"]

    test, train, valid = [], [], []
    for e in entries:
        if e["role"] == "test" or (e["riwaya"] == "warsh" and e["sura"] in d.warsh_test_surahs):
            test.append(e)
        elif rnd.random() < d.valid_fraction:
            valid.append(e)
        else:
            train.append(e)

    # Iqra'Eval: Quranic sentences read by ordinary people (download.py fetches them). Added after
    # the split above, so the everyayah clips keep exactly their train/valid/test assignment.
    if getattr(d, "iqraeval", False):
        iqra = load_iqraeval(cfg)
        if not iqra:
            print("Iqra'Eval: nothing downloaded yet (run download.py first)")
        iqra_durations = load_durations(cfg, [e["path"] for e in iqra]) if iqra else {}
        added = {"train": 0, "valid": 0}
        for e in iqra:
            dur = iqra_durations.get(e["path"])
            if dur is None or dur < 0.1:
                skipped["damaged"] += 1
                continue
            if not (d.min_duration <= dur <= d.max_duration) or len(e["text"]) * 1.3 > dur * 25:
                skipped["duration"] += 1
                continue
            target = "valid" if e["split"] == "dev" else "train"
            (valid if target == "valid" else train).append(dict(
                path=e["path"], duration=round(dur, 3), text=e["text"], reciter="iqraeval", riwaya="hafs",
                sura=e["sura"], aya=e["aya"], source="iqraeval"))
            added[target] += 1
        if iqra:
            print(f"Iqra'Eval: {added['train']} training and {added['valid']} validation clips "
                  f"from ordinary readers")

    # tlog: phone recordings by ordinary people (download.py). Labels checked against the audio;
    # 1 clip in 200 goes to validation. Also after the split, like Iqra'Eval.
    if getattr(d, "tlog_shards", 0):
        tlog = load_tlog(cfg)
        if not tlog:
            print("tlog: nothing downloaded yet (run download.py first)")
        else:
            tlog_durations = load_durations(cfg, [e["path"] for e in tlog])
            usable = []
            for e in tlog:
                dur = tlog_durations.get(e["path"])
                if dur is None or dur < 0.1:
                    skipped["damaged"] += 1
                elif not (d.min_duration <= dur <= min(d.max_duration, 30.0)):
                    skipped["duration"] += 1
                else:
                    usable.append(dict(e, duration=round(dur, 3), riwaya="hafs"))
            heard = transcribe_tlog(cfg, usable)
            limit = getattr(d, "tlog_max_error", 0.35)
            added = {"train": 0, "valid": 0, "wrong label": 0, "relabelled": 0}
            for e in usable:
                err, label = best_label(e, heard[e["path"]], hafs)
                if err > limit:
                    added["wrong label"] += 1
                    continue
                added["relabelled"] += label != e["text"]
                if len(label) * 1.3 > e["duration"] * 25:
                    skipped["too_fast"] += 1
                    continue
                target = "valid" if zlib.crc32(e["file"].encode()) % 200 == 0 else "train"
                (valid if target == "valid" else train).append(dict(
                    path=e["path"], duration=e["duration"], text=label, reciter="tlog", riwaya="hafs",
                    sura=e["sura"], aya=e["aya"], source="tlog"))
                added[target] += 1
            print(f"tlog: {added['train']} training and {added['valid']} validation clips "
                  f"({added['relabelled']} with a corrected label: neighbouring ayah or basmala/isti'adha "
                  f"added); {added['wrong label']} dropped (the audio doesn't match the Quran text around it)")

    out = config_mod.manifest_dir(cfg)
    os.makedirs(out, exist_ok=True)
    for name, rows in (("train", train), ("valid", valid), ("test", test)):
        with open(os.path.join(out, f"{name}.jsonl"), "w", encoding="utf-8", newline="\n") as f:
            for e in rows:
                e = {k: v for k, v in e.items() if k != "role"}
                e["text"] = text.normalize(e["text"])  # lists older versions saved (download.py's)
                f.write(json.dumps(e, ensure_ascii=False) + "\n")
    text.write_tokens(os.path.join(out, "tokens.txt"))

    hours = lambda rows: sum(e["duration"] for e in rows) / 3600  # noqa: E731
    print(f"skipped: {skipped}")
    print(f"train {len(train)} clips / {hours(train):.1f} h   valid {len(valid)} / {hours(valid):.1f} h   "
          f"test {len(test)} / {hours(test):.1f} h")
    per = {}
    for e in train:
        per[(e["riwaya"], e["reciter"])] = per.get((e["riwaya"], e["reciter"]), 0) + e["duration"]
    for (riwaya, folder), s in sorted(per.items()):
        print(f"  {riwaya:5s} {folder:45s} {s / 3600:6.1f} h")
    print("manifests:", out)


if __name__ == "__main__":
    sys.exit(main())