import numpy as np
from dataclasses import dataclass, field
from silero_tts.silero_tts import SileroTTS

from app.segmentation import Segment
from app.ssml import build_ssml_chunks

SSML_MAX_CHARS = 500


@dataclass
class ModelEntry:
    model_id: str
    tts: SileroTTS
    languages: list[str] = field(default_factory=list)


@dataclass
class SynthesisEngine:
    models: list[ModelEntry] = field(default_factory=list)
    sample_rate: int = 24000
    use_ssml: bool = True
    ssml_max_chars: int = SSML_MAX_CHARS

    def load(self, models_config, sample_rate: int = 24000) -> None:
        self.sample_rate = sample_rate
        self.models = []
        for m in models_config:
            tts = SileroTTS(
                model_id=m.file,
                language=m.languages[0] if m.languages else "ru",
                speaker="default",
                sample_rate=sample_rate,
                device="cpu",
                num_threads=4,
            )
            self.models.append(ModelEntry(
                model_id=m.file,
                tts=tts,
                languages=m.languages,
            ))

    def find_model(self, lang: str) -> ModelEntry | None:
        for m in self.models:
            if lang in m.languages:
                return m
        return None

    def get_model(self, model_id: str) -> ModelEntry | None:
        for m in self.models:
            if m.model_id == model_id:
                return m
        return None

    def available_speakers(self) -> dict[str, list[str]]:
        return {m.model_id: list(m.tts.tts_model.speakers) for m in self.models}

    def _synthesize_text(self, model: ModelEntry, speaker: str, text: str) -> np.ndarray:
        audio = model.tts.tts_model.apply_tts(
            text=text,
            speaker=speaker,
            sample_rate=model.tts.sample_rate,
            put_accent=model.tts.put_accent,
            put_yo=model.tts.put_yo,
        )
        return (audio.numpy() * 32767).astype(np.int16)

    def _synthesize_ssml(self, model: ModelEntry, speaker: str, ssml_text: str) -> np.ndarray:
        audio = model.tts.tts_model.apply_tts(
            ssml_text=ssml_text,
            speaker=speaker,
            sample_rate=model.tts.sample_rate,
            put_accent=True,
            put_yo=model.tts.put_yo,
        )
        return (audio.numpy() * 32767).astype(np.int16)

    def synthesize_segments(
        self, segments: list[Segment], speaker: str, lang: str,
        *, rate: str | None = None, pitch: str | None = None, use_ssml: bool | None = None,
    ) -> bytes:
        model = self.find_model(lang)
        if model is None:
            raise ValueError(f"no model supports language '{lang}'")
        return self._synthesize_multi(segments, model, speaker, rate, pitch, use_ssml)

    def synthesize_segments_model(
        self, segments: list[Segment], model_id: str, speaker: str,
        *, rate: str | None = None, pitch: str | None = None, use_ssml: bool | None = None,
    ) -> bytes:
        model = self.get_model(model_id)
        if model is None:
            raise ValueError(f"model '{model_id}' not loaded")
        return self._synthesize_multi(segments, model, speaker, rate, pitch, use_ssml)

    def _synthesize_multi(
        self, segments: list[Segment], model: ModelEntry, speaker: str,
        rate: str | None, pitch: str | None, use_ssml: bool | None,
    ) -> bytes:
        parts: list[np.ndarray] = []
        ssml = self.use_ssml if use_ssml is None else use_ssml
        if ssml:
            for chunk in build_ssml_chunks(
                segments, rate=rate, pitch=pitch, max_chars=self.ssml_max_chars
            ):
                parts.append(self._synthesize_ssml(model, speaker, chunk))
        else:
            for seg in segments:
                parts.append(self._synthesize_text(model, speaker, seg.text))
        if not parts:
            return b""
        return np.concatenate(parts).tobytes()