#!/usr/bin/env python3
"""
Sofelia TTS 82M — minimal inference example (speaker: Eliaa).

    pip install kokoro misaki espeakng-loader phonemizer-fork soundfile torch
    python inference.py "مرحبا كيف حالك اليوم؟"

Arabic has no Kokoro lang_code, so we phonemize with misaki espeak-ng 'ar'
(+ the Sofelia frontend) and call KModel directly, sentence by sentence.
"""

import re
import sys

import numpy as np
import soundfile as sf
import torch
from kokoro import KModel
from misaki import espeak

from sofelia_frontend import text_to_phonemes

MODEL = "kokoro_sofelia_82M.pth"
CONFIG = "config.json"
VOICE = "voices/eliaa.pt"


def main():
    text = sys.argv[1] if len(sys.argv) > 1 else "مرحبا، أنا إيلياء، صوت سُفيليا للعربي الفلسطيني."
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = KModel(repo_id="hexgrad/Kokoro-82M", config=CONFIG, model=MODEL).to(device).eval()
    voice = torch.load(VOICE, map_location="cpu", weights_only=True)
    g2p = espeak.EspeakG2P(language="ar")

    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].to(device), 1.0, return_output=False)
        pieces.append(audio.cpu().numpy().squeeze())
        pieces.append(np.zeros(int(0.25 * 24000), dtype=np.float32))

    out = np.concatenate(pieces[:-1]) if pieces else np.zeros(2400, dtype=np.float32)
    sf.write("output.wav", out, 24000)
    print(f"Saved output.wav ({len(out) / 24000:.1f}s)")


if __name__ == "__main__":
    main()