from typing import Dict, List, Any
import torch
import numpy as np
import base64
import soundfile as sf
import io
from transformers import pipeline

class EndpointHandler:
    def __init__(self, path: str):
        """
        Initialize the endpoint with the model path.
        Args:
            path (str): The file path or model ID for loading the model.
        """
        self.model = pipeline("text-to-speech", model=path)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Process a prediction request using the loaded model.
        Args:
            data (Dict[str, Any]): The request body containing 'inputs' and other parameters.
        Returns:
            List[Dict[str, Any]]: A list containing dictionaries with the model's output.
        """
        inputs = data.get("inputs")
        if not inputs:
            raise ValueError("The 'inputs' key is required in the data dictionary and cannot be empty.")

        if isinstance(inputs, str):
            inputs = [inputs]  # Convert to list to handle consistently as batch

        if not all(isinstance(i, str) for i in inputs):
            raise TypeError("All inputs must be strings.")

        return self.generate_predictions(inputs)

    def generate_predictions(self, texts: List[str]) -> List[Dict[str, Any]]:
        """
        Generate predictions for a list of texts.
        Args:
            texts (List[str]): A list of texts for which to generate predictions.
        Returns:
            Base64 string
        """
        output = self.model(texts[0])
        
        audio_waveform = output["audio"][0]

        buffer = io.BytesIO()
        sf.write(buffer, audio_waveform, output["sampling_rate"], format='WAV')
        buffer.seek(0)  # Rewind the buffer to the beginning
        
        base64_audio = base64.b64encode(buffer.read()).decode('utf-8')
        return base64_audio