#!/usr/bin/env python3
"""
Kokoro Arabic (Palestinian) Training Data Preparation
======================================================
Builds StyleTTS2 train/val lists from two Palestinian Arabic datasets:

  Speaker 0 "Kore"  : /home/ubuntu/data — Gemini TTS story segments
                      (manifests/segments.jsonl + audio/*.wav + text/segments/*.txt)
  Speaker 1 "Sirin" : hamdallah/Sofelia-Palestinian-TTS HF dataset
                      (metadata.csv in LJSpeech format + wavs/*.wav)

Both datasets are already 24kHz mono 16-bit WAV.

G2P: misaki EspeakG2P(language='ar'), with fixups mapping phonemes that are
not in Kokoro-82M's 178-token vocab onto unused-but-real vocab tokens so the
phonemic contrasts of Arabic survive:

  ħ (U+0127 voiceless pharyngeal fric.) -> ʰ (distinct h-like token)
  ʕ (U+0295 voiced pharyngeal fric.)    -> ʁ (voiced uvular fric., unused by ar)
  ˤ (U+02E4 pharyngealization/emphatic) -> ᵊ (distinct modifier token)
  ̪ (U+032A dental diacritic)           -> removed (not phonemic)
  [ ]                                   -> removed (espeak markup)

Usage:
    python3 scripts/prepare_arabic.py
"""

import json
import random
import sys
from multiprocessing import Pool
from pathlib import Path

import soundfile as sf
from tqdm import tqdm

# ── Paths ────────────────────────────────────────────────────────────────────

REPO = Path(__file__).resolve().parent.parent
TRAINING_DIR = REPO / "training"
KOKORO_CONFIG = TRAINING_DIR / "config.json"

KORE_DIR = Path("/home/ubuntu/data")
KORE_MANIFEST = KORE_DIR / "manifests" / "segments.jsonl"

SOFELIA_DIR = Path(
    "/home/ubuntu/.cache/huggingface/hub/datasets--hamdallah--Sofelia-Palestinian-TTS"
    "/snapshots/316fd018d7855ada768639d40b0deb046be89558/dataset"
)

TRAIN_LIST = TRAINING_DIR / "train_list_ar.txt"
VAL_LIST = TRAINING_DIR / "val_list_ar.txt"
OOD_FILE = TRAINING_DIR / "OOD_texts_ar.txt"
STATS_FILE = TRAINING_DIR / "stats_ar.json"

# ── Filters / split ──────────────────────────────────────────────────────────

MIN_DURATION_S = 2.0
MAX_DURATION_S = 20.0  # caps worst-case batch memory (waveform discriminator)
MIN_PHONEME_LEN = 10
VAL_RATIO = 0.02
N_OOD = 128  # entries reserved for OOD texts (excluded from training entirely)
RANDOM_SEED = 42

# ── Phoneme fixups (see module docstring) ────────────────────────────────────

PHONEME_FIXUPS = {
    "ħ": "ʰ",  # ħ -> ʰ
    "ʕ": "ʁ",  # ʕ -> ʁ
    "ˤ": "ᵊ",  # ˤ -> ᵊ
    "̪": "",        # dental diacritic -> drop
    "[": "",
    "]": "",
}

_g2p = None
_marbuta_cache = {}

import json as _json
import re

_MARBUTA_WORD = re.compile(r"([ؠ-يـ]+ة)(?![ؠ-ي])")

# Arabic punctuation → Latin equivalents that espeak passes through to the
# phoneme string (the Arabic forms are silently dropped, so the model never
# sees pause/question anchors without this).
PUNCT_MAP = str.maketrans({"،": ",", "؛": ";", "؟": "?"})

_LEXICON_PATH = REPO / "training" / "ar_lexicon.json"
_lexicon = None
_lexicon_re = None


def _load_lexicon():
    global _lexicon, _lexicon_re
    if _lexicon is None:
        if _LEXICON_PATH.exists():
            entries = _json.loads(_LEXICON_PATH.read_text())
            _lexicon = {k: v for k, v in entries.items() if not k.startswith("_")}
        else:
            _lexicon = {}
        if _lexicon:
            alts = "|".join(re.escape(w) for w in sorted(_lexicon, key=len, reverse=True))
            _lexicon_re = re.compile(rf"(?<![ؠ-يـ])({alts})(?![ؠ-يـ])")
    return _lexicon, _lexicon_re


def normalize_arabic(text, g2p):
    """Full text normalization: punctuation map + pronunciation lexicon + ة fix."""
    text = text.translate(PUNCT_MAP)
    lex, lex_re = _load_lexicon()
    if lex_re is not None:
        text = lex_re.sub(lambda m: lex[m.group(1)], text)
    return normalize_taa_marbuta(text, g2p)


def normalize_taa_marbuta(text, g2p):
    """Fix espeak's OOV fallback that reads word-final ة as /t/ (e.g. قريبة→qriːbt).

    espeak pronounces in-lexicon ة words correctly (مدرسة→a) but renders unknown
    ones with a final t. For each ة-final word, G2P it alone: if the result ends
    in 't', rewrite the trailing ة to a bare fatha (قريبَ), which espeak reads as
    the correct pausal/dialect final 'a' (qriːba). Cached per word type.
    """

    def repl(m):
        word = m.group(1)
        fixed = _marbuta_cache.get(word)
        if fixed is None:
            try:
                ph, _ = g2p(word)
                broken = ph.strip().rstrip(".").endswith("t")
            except Exception:
                broken = False
            fixed = word[:-1] + "َ" if broken else word  # ة → fatha
            _marbuta_cache[word] = fixed
        return fixed

    return _MARBUTA_WORD.sub(repl, text)


def _init_worker():
    global _g2p
    from misaki import espeak

    _g2p = espeak.EspeakG2P(language="ar")


def _phonemize(item):
    """(idx, text) -> (idx, phonemes or None)"""
    idx, text = item
    try:
        text = normalize_arabic(text, _g2p)
        ph, _ = _g2p(text)
    except Exception:
        return idx, None
    for old, new in PHONEME_FIXUPS.items():
        ph = ph.replace(old, new)
    return idx, ph.strip()


def collect_kore():
    """Speaker 0: Gemini story segments."""
    entries = []
    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
    with open(KORE_MANIFEST) as f:
        for line in f:
            e = json.loads(line)
            dur = e.get("durationSeconds", 0)
            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
                dropped["bad_duration"] += 1
                continue
            seg_id = e["segmentId"]
            wav = KORE_DIR / "audio" / f"{seg_id}.wav"
            txt = KORE_DIR / "text" / "segments" / f"{seg_id}.txt"
            if not wav.exists() or not txt.exists():
                dropped["missing_file"] += 1
                continue
            text = txt.read_text().strip()
            if not text:
                dropped["empty_text"] += 1
                continue
            entries.append({"wav": str(wav), "text": text, "speaker": 0, "dur": dur})
    print(f"Kore: kept {len(entries):,}  dropped {dropped}")
    return entries


def collect_sofelia():
    """Speaker 1: Sofelia Palestinian TTS."""
    entries = []
    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
    with open(SOFELIA_DIR / "metadata.csv") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split("|")
            if len(parts) < 2:
                continue
            wav = SOFELIA_DIR / "wavs" / f"{parts[0]}.wav"
            if not wav.exists():
                dropped["missing_file"] += 1
                continue
            text = parts[1].strip()
            if not text:
                dropped["empty_text"] += 1
                continue
            info = sf.info(str(wav))
            dur = info.frames / info.samplerate
            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
                dropped["bad_duration"] += 1
                continue
            entries.append({"wav": str(wav), "text": text, "speaker": 1, "dur": dur})
    print(f"Sofelia: kept {len(entries):,}  dropped {dropped}")
    return entries


def main():
    vocab = json.load(open(KOKORO_CONFIG))["vocab"]

    entries = collect_kore() + collect_sofelia()
    print(f"Total entries before G2P: {len(entries):,}")

    # ── G2P (parallel) ───────────────────────────────────────────────────
    items = [(i, e["text"]) for i, e in enumerate(entries)]
    with Pool(processes=8, initializer=_init_worker) as pool:
        for idx, ph in tqdm(
            pool.imap_unordered(_phonemize, items, chunksize=64),
            total=len(items),
            desc="G2P",
        ):
            entries[idx]["ipa"] = ph

    # ── Validate against Kokoro vocab ────────────────────────────────────
    valid = []
    n_g2p_fail = n_short = n_unknown = 0
    unknown_chars = {}
    for e in entries:
        ph = e.get("ipa")
        if not ph:
            n_g2p_fail += 1
            continue
        if len(ph) < MIN_PHONEME_LEN:
            n_short += 1
            continue
        bad = [ch for ch in ph if ch not in vocab]
        if bad:
            n_unknown += 1
            for ch in bad:
                unknown_chars[ch] = unknown_chars.get(ch, 0) + 1
            continue
        valid.append(e)

    print(f"Valid after G2P: {len(valid):,}")
    print(f"  G2P failures   : {n_g2p_fail}")
    print(f"  Too short      : {n_short}")
    print(f"  Unknown phoneme: {n_unknown}")
    if unknown_chars:
        for ch, c in sorted(unknown_chars.items(), key=lambda x: -x[1])[:20]:
            print(f"    {ch!r} U+{ord(ch):04X}: {c}")

    if not valid:
        print("ERROR: no valid entries")
        sys.exit(1)

    # ── Split: OOD / val / train ─────────────────────────────────────────
    rng = random.Random(RANDOM_SEED)
    rng.shuffle(valid)
    ood_entries = valid[:N_OOD]
    rest = valid[N_OOD:]
    n_val = max(1, int(len(rest) * VAL_RATIO))
    val_entries = rest[:n_val]
    train_entries = rest[n_val:]

    def write_list(path, ents):
        with open(path, "w") as f:
            for e in ents:
                f.write(f"{e['wav']}|{e['ipa']}|{e['speaker']}\n")

    write_list(TRAIN_LIST, train_entries)
    write_list(VAL_LIST, val_entries)

    with open(OOD_FILE, "w") as f:
        f.write("\n".join(e["ipa"] for e in ood_entries) + "\n")

    # ── Stats ────────────────────────────────────────────────────────────
    def hours(ents):
        return round(sum(e["dur"] for e in ents) / 3600, 2)

    stats = {
        "train": {"files": len(train_entries), "hours": hours(train_entries)},
        "val": {"files": len(val_entries), "hours": hours(val_entries)},
        "ood_texts": len(ood_entries),
        "speakers": {
            "0_kore": {
                "files": sum(1 for e in valid if e["speaker"] == 0),
                "hours": hours([e for e in valid if e["speaker"] == 0]),
            },
            "1_sirin": {
                "files": sum(1 for e in valid if e["speaker"] == 1),
                "hours": hours([e for e in valid if e["speaker"] == 1]),
            },
        },
    }
    with open(STATS_FILE, "w") as f:
        json.dump(stats, f, indent=2)

    print(json.dumps(stats, indent=2))
    print(f"\nWrote {TRAIN_LIST}\n      {VAL_LIST}\n      {OOD_FILE}")


if __name__ == "__main__":
    main()