#!/usr/bin/env python3
"""
Kokoro Arabic: Test Inference
==============================
Synthesizes Palestinian Arabic test sentences from a StyleTTS2 checkpoint.

Unlike test_inference.py (German), this bypasses KPipeline: Arabic has no
Kokoro lang_code, so we G2P with misaki espeak-ng 'ar' and apply the same
phoneme fixups as scripts/prepare_arabic.py, then call KModel directly.

Usage:
    CUDA_VISIBLE_DEVICES= python3 scripts/test_inference_arabic.py \
        --checkpoint StyleTTS2/logs/kokoro-arabic/epoch_1st_00000.pth \
        --voicepack voices/ar_kore_e1.pt \
        --output-dir test_output/epoch1_ar --device cpu
"""

import argparse
import sys
from pathlib import Path

_repo_root = Path(__file__).resolve().parents[1]
_kokoro_submodule = _repo_root / "kokoro"
if _kokoro_submodule.exists() and str(_kokoro_submodule) not in sys.path:
    sys.path.insert(0, str(_kokoro_submodule))

sys.path.insert(0, str(_repo_root / "scripts"))
from prepare_arabic import PHONEME_FIXUPS, normalize_arabic  # noqa: E402
from test_inference import convert_checkpoint  # noqa: E402

TEST_SENTENCES = [
    # Daily-life Palestinian dialect
    "مرحبا، كيف حالك اليوم؟ إن شاء الله تكون مبسوط.",
    "بدي أروح ع السوق أشتري خضرة وفواكه للبيت.",
    # Pharyngeals and emphatics (ع، ح، ص، ض، ط)
    "الصبح الساعة سبعة طلعت من البيت على الشغل.",
    "حكيت مع عمي عن العرس اللي صار الأسبوع الماضي.",
    # Numbers and prices
    "هاد حقه خمسين شيكل بس، مش غالي.",
    # Question prosody
    "شو رأيك نطلع نتمشى شوي بعد العصر؟",
    # Qaf-heavy (ق realization check)
    "قلت لقاسم القهوة بالسوق قريبة من البنك القديم.",
]


def main():
    parser = argparse.ArgumentParser(description="Test Arabic Kokoro inference")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--checkpoint", help="StyleTTS2 checkpoint (.pth), converted automatically")
    group.add_argument("--model", help="Already-converted Kokoro-format weights (.pth)")
    parser.add_argument("--voicepack", required=True)
    parser.add_argument("--config", default="training/config.json")
    parser.add_argument("--output-dir", default="test_output/arabic")
    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    parser.add_argument("--speed", type=float, default=1.0)
    args = parser.parse_args()

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

    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)

    if args.checkpoint:
        model_path = convert_checkpoint(
            args.checkpoint, str(out / "kokoro_arabic_converted.pth")
        )
    else:
        model_path = args.model

    print(f"Loading KModel on {args.device}...")
    kmodel = KModel(
        repo_id="hexgrad/Kokoro-82M", config=args.config, model=model_path
    )
    kmodel = kmodel.to(args.device).eval()

    voice = torch.load(args.voicepack, map_location="cpu", weights_only=True)
    print(f"Voicepack: {args.voicepack} {tuple(voice.shape)}")

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

    for i, text in enumerate(TEST_SENTENCES):
        print(f"\n[{i + 1}/{len(TEST_SENTENCES)}] {text}")
        ps, _ = g2p(normalize_arabic(text, g2p))
        for old, new in PHONEME_FIXUPS.items():
            ps = ps.replace(old, new)
        ps = ps.strip()
        print(f"  phonemes: {ps}")
        if len(ps) > 510:
            ps = ps[:510]
        ref_s = voice[len(ps) - 1]
        with torch.no_grad():
            audio = kmodel(ps, ref_s, args.speed, return_output=False)
        audio = audio.cpu().numpy().squeeze()
        wav_path = out / f"ar_test_{i + 1:02d}.wav"
        sf.write(str(wav_path), audio, 24000)
        print(f"  saved: {wav_path} ({len(audio) / 24000:.1f}s)")

    print(f"\nDone! Output in {out}/")


if __name__ == "__main__":
    main()