"""
FroxAI Flex-Audio — inference helper.

Loads the fine-tuned XTTS-v2 checkpoint (model.pth + config.json + vocab.json,
all expected in the repo root) and exposes a single generate() function used by
both app.py (Gradio UI) and any other script that wants to import this directly.
"""

import os
import torch
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
VOICES_DIR = os.path.join(REPO_ROOT, "voice_refs")

_model = None
_config = None


def load_model():
    """Loads the model once and caches it. Safe to call repeatedly."""
    global _model, _config
    if _model is not None:
        return _model, _config

    config = XttsConfig()
    config.load_json(os.path.join(REPO_ROOT, "config.json"))

    model = Xtts.init_from_config(config)
    model.load_checkpoint(config, checkpoint_dir=REPO_ROOT, eval=True)
    if torch.cuda.is_available():
        model.cuda()

    _model, _config = model, config
    return model, config


def list_voices():
    """Returns available voice names (without the .wav extension) found in voice_refs/."""
    if not os.path.isdir(VOICES_DIR):
        return []
    return sorted(
        os.path.splitext(f)[0]
        for f in os.listdir(VOICES_DIR)
        if f.lower().endswith(".wav")
    )


def generate(text: str, language: str, voice: str) -> str:
    """Generates speech and returns the path to the written WAV file.

    Args:
        text: the text to speak.
        language: an XTTS-v2 language code, e.g. 'en', 'es', 'hi'.
        voice: a voice name as returned by list_voices() (matches a file in voice_refs/).
    """
    model, config = load_model()

    voice_path = os.path.join(VOICES_DIR, f"{voice}.wav")
    if not os.path.exists(voice_path):
        raise FileNotFoundError(f"No reference clip found for voice '{voice}' at {voice_path}")

    outputs = model.synthesize(
        text,
        config,
        speaker_wav=voice_path,
        language=language,
    )

    out_path = "/tmp/frox_flex_audio_output.wav"
    import soundfile as sf
    sf.write(out_path, outputs["wav"], config.audio.sample_rate)
    return out_path