"""
Zero-shot voice cloning engine.

Extracted and simplified from jamiepine/voicebox's Qwen3-TTS "Base" backend
(backend/backends/pytorch_backend.py). Only the voice-cloning path is kept:
give a short reference clip + its transcript, then synthesize any text in
that cloned voice.

Model: Qwen/Qwen3-TTS-12Hz-{1.7B,0.6B}-Base via the `qwen-tts` package.
Weights download automatically from HuggingFace on first use.
"""

from __future__ import annotations

import logging
import threading
from typing import Optional, Tuple

import numpy as np

logger = logging.getLogger("voice_copy.engine")

# Qwen3-TTS speaks these languages. "auto" lets the model detect it.
LANGUAGE_CODE_TO_NAME = {
    "auto": "auto",
    "ko": "korean",
    "en": "english",
    "zh": "chinese",
    "ja": "japanese",
    "de": "german",
    "fr": "french",
    "ru": "russian",
    "pt": "portuguese",
    "es": "spanish",
    "it": "italian",
}

HF_MODEL_MAP = {
    "1.7B": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
    "0.6B": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
}


def pick_device() -> str:
    """Best available torch device: CUDA GPU if present, else CPU."""
    try:
        import torch
    except Exception:
        return "cpu"
    if torch.cuda.is_available():
        return "cuda"
    # Apple Silicon
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class VoiceCloner:
    """Lazily-loaded Qwen3-TTS voice-clone wrapper.

    The heavy model is only loaded on the first `clone()` call, so the web
    server can start instantly and the (multi-GB) download happens on demand.
    """

    def __init__(self, model_size: str = "1.7B", device: Optional[str] = None):
        if model_size not in HF_MODEL_MAP:
            raise ValueError(f"Unknown model size {model_size!r}; choose one of {list(HF_MODEL_MAP)}")
        self.model = None
        self.model_size = model_size
        self.device = device or pick_device()
        self._lock = threading.Lock()

    @property
    def loaded(self) -> bool:
        return self.model is not None

    def load(self) -> None:
        """Load the model (thread-safe, idempotent)."""
        if self.model is not None:
            return
        with self._lock:
            if self.model is not None:
                return
            import torch
            from qwen_tts import Qwen3TTSModel

            model_path = HF_MODEL_MAP[self.model_size]
            logger.info("Loading %s on %s (first run downloads weights)...", model_path, self.device)

            if self.device == "cpu":
                self.model = Qwen3TTSModel.from_pretrained(
                    model_path,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=False,
                )
            else:
                self.model = Qwen3TTSModel.from_pretrained(
                    model_path,
                    device_map=self.device,
                    torch_dtype=torch.bfloat16,
                )
            logger.info("Model loaded.")

    def clone(
        self,
        ref_audio_path: str,
        ref_text: str,
        text: str,
        language: str = "auto",
        instruct: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> Tuple[np.ndarray, int]:
        """Synthesize `text` in the voice of `ref_audio_path`.

        Args:
            ref_audio_path: Path to a short reference clip (a few seconds is enough).
            ref_text: Exact transcript of the reference clip.
            text: The text to speak in the cloned voice.
            language: Language code from LANGUAGE_CODE_TO_NAME ("auto" to detect).
            instruct: Optional natural-language style hint (e.g. "speak cheerfully").
            seed: Optional random seed for reproducible output.

        Returns:
            (audio_float32_mono, sample_rate)
        """
        self.load()
        import torch

        if seed is not None:
            torch.manual_seed(seed)
            if self.device == "cuda":
                torch.cuda.manual_seed_all(seed)

        # Build the speaker prompt from the reference clip + its transcript.
        voice_prompt = self.model.create_voice_clone_prompt(
            ref_audio=str(ref_audio_path),
            ref_text=ref_text,
            x_vector_only_mode=False,
        )

        wavs, sample_rate = self.model.generate_voice_clone(
            text=text,
            voice_clone_prompt=voice_prompt,
            language=LANGUAGE_CODE_TO_NAME.get(language, "auto"),
            instruct=instruct or None,
        )

        audio = wavs[0]
        if isinstance(audio, torch.Tensor):
            audio = audio.squeeze().cpu().numpy()
        return np.asarray(audio, dtype=np.float32), int(sample_rate)