from pathlib import Path
from typing import Any, Dict, List
import sys
import numpy as np

import base64
import io
import soundfile as sf


REPO_ROOT = Path(__file__).resolve().parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from cosyvoice.cli.cosyvoice import CosyVoice

class EndpointHandler:
    def __init__(self, model_dir: str):
        self.model_dir = model_dir  # "/repository"
        cosy_model_dir = str(Path(self.model_dir) / "model")  # "/repository/model"

        self.cosyvoice = CosyVoice(
            cosy_model_dir,
            load_jit=False,
            load_trt=False,
            fp16=False,
        )


    def _b64_to_bytes(self, s: str) -> bytes:
        return base64.b64decode(s)

    def _bytes_to_b64_wav(self, audio: np.ndarray, sr: int) -> str:
        buf = io.BytesIO()
        sf.write(buf, audio.T, sr, format="WAV")
        return base64.b64encode(buf.getvalue()).decode("utf-8")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Expected payload:
        {
          "inputs": {
            "text": "...",
            "ref_text": "...",
            "ref_audio_b64": "<base64 wav>"
          }
        }
        """
        inputs = data["inputs"]
        text = inputs["text"]
        ref_text = inputs.get("ref_text", "")
        ref_audio_b64 = inputs["ref_audio_b64"]

        ref_bytes = self._b64_to_bytes(ref_audio_b64)
        ref_path = "/tmp/ref.wav"
        with open(ref_path, "wb") as f:
            f.write(ref_bytes)

        chunks = []
        for chunk in self.cosyvoice.inference_zero_shot(
            text,
            ref_text,
            ref_path,
            stream=False,
        ):
            audio = chunk.get("tts_speech", chunk.get("audio"))
            chunks.append(audio)

        audio = np.concatenate(chunks, axis=-1)
        sr = self.cosyvoice.sample_rate
        audio_b64 = self._bytes_to_b64_wav(audio, sr)

        return [{"audio_wav_base64": audio_b64, "sample_rate": sr}]