"""Persian text normalization for TTS: digits to Persian words, long-text chunking.

The model reads Persian words, not digits — no TTS reads "۱۹۸۱" aloud. This
front-end rewrites numbers into Persian script words BEFORE synthesis, and
splits long text into sentence-sized chunks (training capped utterances at
12 s; walls of text degrade everything).

    normalize("سال ۱۹۸۱")  -> "سال هزار و نهصد و هشتاد و یک"
    chunk(long_text)        -> ["sentence 1.", "sentence 2.", ...]

Pure stdlib. Run directly for the self-check:  python scripts/fa_normalize.py
"""
import re

ONES = ["", "یک", "دو", "سه", "چهار", "پنج", "شش", "هفت", "هشت", "نه"]
TEENS = ["ده", "یازده", "دوازده", "سیزده", "چهارده", "پانزده",
         "شانزده", "هفده", "هجده", "نوزده"]
TENS = ["", "", "بیست", "سی", "چهل", "پنجاه", "شصت", "هفتاد", "هشتاد", "نود"]
HUNDREDS = ["", "صد", "دویست", "سیصد", "چهارصد", "پانصد",
            "ششصد", "هفتصد", "هشتصد", "نهصد"]
SCALES = ["", "هزار", "میلیون", "میلیارد", "تریلیون"]

_DIGIT_MAP = str.maketrans("۰۱۲۳۴۵۶۷۸۹٠١٢٣٤٥٦٧٨٩", "01234567890123456789")


def _under_1000(n: int) -> str:
    parts = []
    if n >= 100:
        parts.append(HUNDREDS[n // 100])
        n %= 100
    if 10 <= n <= 19:
        parts.append(TEENS[n - 10])
    else:
        if n >= 20:
            parts.append(TENS[n // 10])
            n %= 10
        if n:
            parts.append(ONES[n])
    return " و ".join(parts)


def number_to_persian(n: int) -> str:
    if n == 0:
        return "صفر"
    if n < 0:
        return "منفی " + number_to_persian(-n)
    groups = []
    scale = 0
    while n:
        n, g = divmod(n, 1000)
        if g:
            words = _under_1000(g)
            # "یک هزار" is unidiomatic; plain "هزار" is what people say
            if scale == 1 and g == 1:
                words = ""
            groups.append((words + " " + SCALES[scale]).strip())
        scale += 1
    return " و ".join(reversed(groups))


def _num_repl(m: re.Match) -> str:
    intpart, _, frac = m.group(0).partition("٫")
    n = int(intpart.replace(",", "").replace("٬", ""))
    out = number_to_persian(n)
    if frac:
        out += " ممیز " + " ".join(ONES[int(d)] if d != "0" else "صفر" for d in frac)
    return out


# integers with optional thousands separators and optional decimal part (٫)
_NUM_RE = re.compile(r"\d[\d,٬]*(?:٫\d+)?")
_PERCENT_RE = re.compile(r"([\d,٬٫]+)\s*[٪%]")


def normalize(text: str) -> str:
    text = text.translate(_DIGIT_MAP)
    text = _PERCENT_RE.sub(lambda m: m.group(1) + " درصد", text)
    return _NUM_RE.sub(_num_repl, text)


_SENT_SPLIT = re.compile(r"(?<=[.!?؟؛…])\s+|\n+")
MAX_CHUNK = 220  # chars ≈ a 10-12 s utterance, the training cap


def chunk(text: str) -> list[str]:
    """Sentence-split; greedily pack sentences into chunks under MAX_CHUNK."""
    sentences = [s.strip() for s in _SENT_SPLIT.split(text) if s.strip()]
    chunks, cur = [], ""
    for s in sentences:
        while len(s) > MAX_CHUNK:  # single overlong sentence: cut at a comma/space
            cut = max(s.rfind("،", 0, MAX_CHUNK), s.rfind(" ", 0, MAX_CHUNK))
            cut = cut if cut > 0 else MAX_CHUNK
            piece, s = s[:cut].strip(), s[cut:].lstrip("، ").strip()
            chunks.append((cur + " " + piece).strip() if cur else piece)
            cur = ""
        if cur and len(cur) + len(s) + 1 > MAX_CHUNK:
            chunks.append(cur)
            cur = s
        else:
            cur = f"{cur} {s}".strip()
    if cur:
        chunks.append(cur)
    return chunks


if __name__ == "__main__":
    assert number_to_persian(1981) == "هزار و نهصد و هشتاد و یک"
    assert number_to_persian(17) == "هفده"
    assert number_to_persian(0) == "صفر"
    assert number_to_persian(2450000) == "دو میلیون و چهارصد و پنجاه هزار"
    assert normalize("سال ۱۹۸۱") == "سال هزار و نهصد و هشتاد و یک"
    assert normalize("۱۷ ساله") == "هفده ساله"
    assert normalize("۹۵٪") == "نود و پنج درصد"
    para = ("داستان در لس آنجلس سال ۱۹۸۱ اتفاق می‌افتد و ماجرای نسخه ۱۷ ساله‌ای از "
            "برت ایستون الیس را در آخرین سال تحصیلش در مدرسه‌ی سطح بالای «باکلی» دنبال می‌کند. "
            "با ورود یک دانش‌آموز جدید و مرموز به نام رابرت مالوری، دنیای او زیر و رو می‌شود؛ "
            "حضوری نگران‌کننده که با فعالیت‌های یک قاتل زنجیره‌ای معروف به «تراولر» همزمان شده است.")
    out = chunk(normalize(para))
    assert not re.search(r"[\d۰-۹]", " ".join(out)), "digits survived"
    assert all(len(c) <= MAX_CHUNK for c in out)
    assert len(out) >= 2
    print(f"self-check OK; paragraph -> {len(out)} chunks:")
    for c in out:
        print(" •", c[:80], "..." if len(c) > 80 else "")