from __future__ import annotations

import threading
from functools import lru_cache
from pathlib import Path

import soundfile as sf
import torch
from pocket_tts import TTSModel

_INFERENCE_LOCK = threading.Lock()
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


@lru_cache(maxsize=1)
def _load_model() -> TTSModel:
    model = TTSModel.load_model(language="english")
    model.to(DEVICE)
    model.eval()
    return model


@lru_cache(maxsize=12)
def _load_voice_state(voice: str):
    return _load_model().get_state_for_audio_prompt(voice)


def load_model() -> None:
    """Load the model now, so that the first request does not wait on it."""
    _load_model()


def synthesize(text: str, voice: str, output_path: Path) -> Path:
    model = _load_model()
    voice_state = _load_voice_state(voice)
    # generate_audio is not thread-safe.
    with _INFERENCE_LOCK:
        audio = model.generate_audio(voice_state, text, copy_state=True)
    waveform = audio.detach().cpu().float().squeeze().numpy()
    sf.write(str(output_path), waveform, model.sample_rate, subtype="PCM_16")
    return output_path