import base64
import io
import sys
import tempfile
from pathlib import Path

import torch
import torchaudio

REPO_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_DIR))
sys.path.insert(0, str(REPO_DIR / "third_party" / "Matcha-TTS"))

from accent_config import ACCENT_INSTRUCTIONS  # noqa: E402


class EndpointHandler:
    def __init__(self, path=""):
        self.model_dir = Path(path or __file__).resolve()
        if self.model_dir.is_file():
            self.model_dir = self.model_dir.parent
        from cosyvoice.cli.cosyvoice import AutoModel

        self.model = AutoModel(
            model_dir=str(self.model_dir),
            fp16=torch.cuda.is_available(),
            load_vllm=False,
        )

    def __call__(self, data):
        text = data.get("inputs", "")
        parameters = data.get("parameters", {})
        accent = parameters.get("accent", "singapore")
        if not text:
            raise ValueError("inputs must contain Chinese text")
        if accent not in ACCENT_INSTRUCTIONS:
            raise ValueError(f"accent must be one of: {', '.join(ACCENT_INSTRUCTIONS)}")

        prompt_path = self.model_dir / "zero_shot_prompt.wav"
        prompt_audio = parameters.get("prompt_audio_base64")
        temp_path = None
        if prompt_audio:
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
                handle.write(base64.b64decode(prompt_audio))
                temp_path = Path(handle.name)
                prompt_path = temp_path
        try:
            chunks = [item["tts_speech"] for item in self.model.inference_instruct2(
                text,
                ACCENT_INSTRUCTIONS[accent],
                str(prompt_path),
                stream=False,
                speed=float(parameters.get("speed", 1.0)),
            )]
            speech = torch.cat(chunks, dim=1).cpu()
            buffer = io.BytesIO()
            torchaudio.save(buffer, speech, self.model.sample_rate, format="wav")
            return {
                "audio_base64": base64.b64encode(buffer.getvalue()).decode("ascii"),
                "sample_rate": self.model.sample_rate,
                "accent": accent,
            }
        finally:
            if temp_path is not None:
                temp_path.unlink(missing_ok=True)