#!/usr/bin/env python3
"""
Synthesize Tamil-script text with the female VITS checkpoint.

Usage:
    python inference.py "சொராஷ்ட்ர மொழி" -o output.wav
    python inference.py "சொராஷ்ட்ர மொழி" --gpu 1
    python inference.py "சொராஷ்ட்ர மொழி" --cpu
"""
import argparse
import os
import re
import unicodedata
from pathlib import Path

import torch
from TTS.api import TTS

CHECKPOINT = "best_model.pth"
MODEL_DIR = Path(__file__).parent


def clean_text(text: str) -> str:
    text = re.sub(r"(?:\{[^}]*\}\s*)+", lambda m: re.findall(r"\{([^}]*)\}", m.group(0))[-1].strip(), text)
    text = unicodedata.normalize("NFC", text)
    text = text.replace(":", "").replace(".", "").replace("'", "").replace("’", "")
    return re.sub(r"  +", " ", text).strip()


def synthesize(tts, text: str, output_path: str) -> None:
    tts.tts_to_file(text=text, file_path=output_path)


def main():
    parser = argparse.ArgumentParser(description="Female Tamil-script TTS inference")
    parser.add_argument("text", help="Input text in Tamil script")
    parser.add_argument("-o", "--output", default="output.wav", help="Output WAV path (default: output.wav)")
    parser.add_argument("--gpu", type=int, default=0, help="GPU device ID (default: 0)")
    parser.add_argument("--cpu", action="store_true", help="Force CPU inference")
    args = parser.parse_args()

    if not args.cpu:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

    device = "cpu" if args.cpu else ("cuda:0" if torch.cuda.is_available() else "cpu")

    tts = TTS(
        model_path=str(MODEL_DIR / CHECKPOINT),
        config_path=str(MODEL_DIR / "config.json"),
        progress_bar=False,
    )
    tts.to(device)

    text = clean_text(args.text)
    print(f"Device    : {device}")
    print(f"Text      : {text}")
    print(f"Output    : {args.output}")

    synthesize(tts, text, args.output)
    print("Done.")


if __name__ == "__main__":
    main()