import math
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import soundfile as sf

from backend.types import VoiceConfig
from backend.voice_presets import VOICE_PRESET_BY_ID

try:
    import spaces
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(fn=None, **_kwargs):
            def decorate(inner):
                return inner

            if fn is not None:
                return decorate(fn)
            return decorate

    spaces = _SpacesShim()

class OmniVoiceAdapter:
    def __init__(
        self,
        model_id: str = "k2-fsa/OmniVoice",
        device_map: str = "cuda:0",
        dtype_name: str = "float16",
    ) -> None:
        self.model_id = model_id
        self.device_map = device_map
        self.dtype_name = dtype_name
        self._model = None
        self._backend = "fallback"
        self._voice_clone_prompt_cache = {}

    def _load_model(self):
        if self._model is not None:
            return self._model
        try:
            import torch
            from omnivoice import OmniVoice
        except Exception:
            self._backend = "fallback"
            return None

        dtype = getattr(torch, self.dtype_name)
        self._model = OmniVoice.from_pretrained(
            self.model_id,
            device_map=self.device_map,
            dtype=dtype,
        )
        self._backend = "omnivoice"
        return self._model

    @spaces.GPU(duration=300)
    def synthesize(
        self,
        *,
        text: str,
        output_path: Path,
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
        language: Optional[str] = None,
    ) -> Dict[str, object]:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        model = self._load_model()
        if model is None:
            return self._synthesize_fallback(
                text=text,
                output_path=output_path,
                voice_config=voice_config,
                speed=speed,
            )

        from omnivoice.models.omnivoice import OmniVoiceGenerationConfig

        kwargs = {
            "text": text,
            "speed": speed,
            "generation_config": OmniVoiceGenerationConfig(
                num_step=diffusion_steps,
                position_temperature=0.0,
                class_temperature=0.0,
            ),
        }
        if language:
            kwargs["language"] = language

        if voice_config.mode == "clone":
            kwargs["voice_clone_prompt"] = self._voice_clone_prompt(voice_config, model)
        elif voice_config.mode == "design":
            kwargs["instruct"] = voice_config.design_prompt
        elif voice_config.narrator_id:
            preset = VOICE_PRESET_BY_ID.get(voice_config.narrator_id)
            if preset:
                kwargs["instruct"] = preset["instruct"]

        audio = model.generate(**kwargs)
        waveform = np.asarray(audio[0], dtype=np.float32)
        sample_rate = 24000
        sf.write(str(output_path), waveform, sample_rate)
        duration_seconds = int(round(len(waveform) / sample_rate))
        return {
            "duration_seconds": max(1, duration_seconds),
            "sample_rate": sample_rate,
            "backend": "local",
            "model": "omnivoice",
            "engine": self._backend,
        }

    def _voice_clone_prompt(self, voice_config: VoiceConfig, model):
        cache_key = (
            voice_config.sample_path or "",
            voice_config.reference_text or "",
        )
        prompt = self._voice_clone_prompt_cache.get(cache_key)
        if prompt is not None:
            return prompt
        prompt = model.create_voice_clone_prompt(
            ref_audio=voice_config.sample_path,
            ref_text=voice_config.reference_text,
        )
        self._voice_clone_prompt_cache[cache_key] = prompt
        return prompt

    def _synthesize_fallback(
        self,
        *,
        text: str,
        output_path: Path,
        voice_config: VoiceConfig,
        speed: float,
    ) -> Dict[str, object]:
        sample_rate = 24000
        duration_seconds = max(1.0, min(20.0, len(text.split()) / max(speed, 0.5) * 0.45))
        total_samples = int(sample_rate * duration_seconds)
        base_freq = 180.0
        if voice_config.mode == "design":
            base_freq = 220.0
        elif voice_config.mode == "clone":
            base_freq = 140.0
        elif voice_config.narrator_id:
            preset_ids = list(VOICE_PRESET_BY_ID.keys())
            if voice_config.narrator_id in preset_ids:
                base_freq = 160.0 + (preset_ids.index(voice_config.narrator_id) * 20)

        timeline = np.linspace(0, duration_seconds, total_samples, endpoint=False)
        waveform = (
            0.15 * np.sin(2 * math.pi * base_freq * timeline)
            + 0.05 * np.sin(2 * math.pi * (base_freq / 2.0) * timeline)
        ).astype(np.float32)
        sf.write(str(output_path), waveform, sample_rate)
        return {
            "duration_seconds": int(round(duration_seconds)),
            "sample_rate": sample_rate,
            "backend": "local",
            "model": "omnivoice",
            "engine": self._backend,
        }