#!/usr/bin/env python3
"""Sofelia TTS — Palestinian Arabic (speaker Eliaa). CPU Gradio Space."""

import re

import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from kokoro import KModel
from misaki import espeak

from sofelia_frontend import text_to_phonemes

REPO = "hamdallah/Sofelia-TTS-82M"
SR = 24000
torch.set_num_threads(4)

print("Downloading model files…")
MODEL_PATH = hf_hub_download(REPO, "kokoro_sofelia_82M.pth")
CONFIG_PATH = hf_hub_download(REPO, "config.json")
VOICE_PATH = hf_hub_download(REPO, "voices/eliaa.pt")

import kokoro as _kokoro_pkg

print(f"[diag] kokoro from: {_kokoro_pkg.__file__}", flush=True)

print("Loading model on CPU…")
MODEL = KModel(repo_id="hexgrad/Kokoro-82M", config=CONFIG_PATH, model=MODEL_PATH).to("cpu").eval()
VOICE = torch.load(VOICE_PATH, map_location="cpu", weights_only=True)
G2P = espeak.EspeakG2P(language="ar")

# ── startup self-test: log phonemes + audio stats so noise vs speech is visible
try:
    import numpy as _np
    _t = "شو أخبارك؟ كلشي تمام إن شاء الله."
    _ps = text_to_phonemes(_t, G2P)[:510]
    print(f"[diag] phonemes: {_ps}", flush=True)
    with torch.no_grad():
        _a = MODEL(_ps, VOICE[len(_ps) - 1], 1.0, return_output=False).cpu().numpy().squeeze()
    print(
        f"[diag] selftest rms={float(_np.sqrt((_a**2).mean())):.4f} "
        f"peak={float(_np.abs(_a).max()):.3f} len={len(_a)/24000:.1f}s "
        f"(clean ref: rms~0.104 peak~0.70 len~3.6s)",
        flush=True,
    )
except Exception as _e:
    print(f"[diag] selftest failed: {_e}", flush=True)

EXAMPLES = [
    "مرحبا، أنا إيلياء. كيف بقدر أساعدك اليوم؟",
    "بدي أروح ع السوق أشتري خضرة وفواكه للبيت.",
    "يا زلمة وين كنت مبارح؟ دورت عليك وما لقيتك.",
    "الصبح الساعة سبعة طلعت من البيت على الشغل.",
]


def compress_silence(x, thresh=0.012, max_gap=0.22, keep=0.14):
    """Cap any silent run longer than max_gap down to `keep` seconds.
    Fixes the ~1s pauses at '.'/',' (model's own silence + chunk-stitch silence)."""
    sil = np.abs(x) < thresh
    keep_n, max_n = int(keep * SR), int(max_gap * SR)
    out, i, n = [], 0, len(x)
    while i < n:
        j = i
        if sil[i]:
            while j < n and sil[j]:
                j += 1
            out.append(x[i : i + keep_n] if (j - i) > max_n else x[i:j])
        else:
            while j < n and not sil[j]:
                j += 1
            out.append(x[i:j])
        i = j
    return np.concatenate(out) if out else x


def synth(text, speed):
    text = (text or "").strip()
    if not text:
        return None
    chunks = [c.strip() for c in re.split(r"(?<=[.!؟?:])\s+", text) if c.strip()] or [text]
    pieces = []
    for ch in chunks:
        ps = text_to_phonemes(ch, G2P)[:510]
        if not ps:
            continue
        with torch.no_grad():
            audio = MODEL(ps, VOICE[len(ps) - 1], float(speed), return_output=False)
        pieces.append(audio.cpu().numpy().squeeze().astype(np.float32))
    if not pieces:
        return None
    full = compress_silence(np.concatenate(pieces))
    # return 16-bit PCM, not float — gradio's float-audio path distorts/clips
    pcm = (np.clip(full, -1.0, 1.0) * 32767.0).astype(np.int16)
    return (SR, pcm)


with gr.Blocks(title="Sofelia TTS — Palestinian Arabic") as demo:
    gr.Markdown(
        "# 🗣️ Sofelia TTS 82M — Palestinian Arabic\n"
        "Natural Palestinian Arabic text-to-speech. Voice: **Eliaa**. Runs on CPU.\n\n"
        "Type Arabic text (dialect welcome) and press Generate."
    )
    with gr.Row():
        with gr.Column():
            txt = gr.Textbox(label="النص العربي", rtl=True, lines=4, value=EXAMPLES[0])
            speed = gr.Slider(0.7, 1.3, value=1.0, step=0.05, label="Speed / السرعة")
            btn = gr.Button("🔊 توليد الصوت", variant="primary")
        out = gr.Audio(label="الصوت", type="numpy")
    gr.Examples(EXAMPLES, inputs=txt)
    btn.click(synth, [txt, speed], out)

if __name__ == "__main__":
    demo.launch()