import requests
import zipfile
import io
import subprocess
import os 
import shutil
from TTS.utils.synthesizer import Synthesizer
from IndicTTS.inference.src.inference import TextToSpeechEngine
import torch

class EndpointHandler():
    def __init__(self,path=""):
        self.path = path
        print(f"python version>> ",{torch.__version__})
        print(f"cuda version>> ",{torch.cuda.get_device_capability(0)[0]})
        models = {}
        odia_model = Synthesizer(
            tts_checkpoint="/repository/or/fastpitch/best_model.pth",
            tts_config_path="/repository/or/fastpitch/config.json",
            # modify this config.json to proper model path
            tts_speakers_file="/repository/or/fastpitch/speakers.pth",
            vocoder_checkpoint="/repository/or/hifigan/best_model.pth",
            vocoder_config="/repository/or/hifigan/config.json",
            use_cuda = True 
        )
        models["or"] = odia_model
        self.engine = TextToSpeechEngine(models)

    def __call__(self, data:dict):
        input_text = data["inputs"]["input_text"]
        input_lang_code = data["inputs"]["input_lang_code"]
        output_lang_code = data["inputs"]["output_lang_code"]
        speaker_gender = data["inputs"]["speaker_gender"]
        audio_array= self.engine.infer_from_text(
        input_text=input_text,
        lang=output_lang_code,
        speaker_name=speaker_gender,
        )
        return [{"audio_array": audio_array}]