import torch
import torchaudio
import io
import base64
from typing import Dict, Any
from chatterbox.mtl_tts import ChatterboxMultilingualTTS


class EndpointHandler:
    def __init__(self, path=""):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"🚀 Loading Chatterbox-Multilingual model on {device}...")
        self.model = ChatterboxMultilingualTTS.from_pretrained(device=device)
        self.device = device
        print("✅ Model loaded successfully")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.pop("inputs", data)
        
        text = inputs.get("text")
        if not text:
            return {"error": "Missing required parameter: text"}
        
        reference_audio = inputs.get("reference_audio")
        if not reference_audio:
            return {"error": "Missing required parameter: reference_audio"}
        
        cfg_weight = inputs.get("cfg_weight", 1.0)
        exaggeration = inputs.get("exaggeration", 1.0)
        language_id = inputs.get("language", "en")
        
        if isinstance(reference_audio, str) and reference_audio.startswith("data:"):
            header, encoded = reference_audio.split(",", 1)
            reference_audio_bytes = base64.b64decode(encoded)
        elif isinstance(reference_audio, str):
            reference_audio_bytes = base64.b64decode(reference_audio)
        else:
            reference_audio_bytes = reference_audio
        
        with open("/tmp/reference_audio.tmp", "wb") as f:
            f.write(reference_audio_bytes)
        
        try:
            wav = self.model.generate(
                text,
                audio_prompt_path="/tmp/reference_audio.tmp",
                language_id=language_id,
                exaggeration=exaggeration,
                cfg_weight=cfg_weight
            )
            
            buffer = io.BytesIO()
            torchaudio.save(buffer, wav, self.model.sr, format="wav")
            buffer.seek(0)
            audio_base64 = base64.b64encode(buffer.read()).decode("utf-8")
            
            return {
                "audio": audio_base64,
                "sample_rate": self.model.sr,
                "format": "wav"
            }
        except Exception as e:
            return {"error": str(e)}