from transcribe.pipelines import WhisperPipe, TranslatePipe, MetaItem, VadPipe
import multiprocessing  as mp
import config

class TranslatePipes:
    def __init__(self) -> None:
        
        # self.whisper_input_q = mp.Queue() 
        # self.translate_input_q = mp.Queue()
        # self.result_queue = mp.Queue()
        
        # whisper 转录
        self._whisper_pipe = self._launch_process(WhisperPipe())
        
        # llm 翻译
        self._translate_pipe = self._launch_process(TranslatePipe())

        # vad 
        self._vad_pipe = self._launch_process(VadPipe())
    
    def _launch_process(self, process_obj):
        process_obj.daemon = True
        process_obj.start()
        return process_obj

    def wait_ready(self):
        self._whisper_pipe.wait()
        self._translate_pipe.wait()
        self._vad_pipe.wait()
    
    def translate(self, text, src_lang, dst_lang) -> MetaItem:
        item = MetaItem(
            transcribe_content=text,
            source_language=src_lang, 
            destination_language=dst_lang)
        self._translate_pipe.input_queue.put(item)
        return self._translate_pipe.output_queue.get()
    

    def transcrible(self, audio_buffer:bytes, src_lang: str) -> MetaItem:
        item = MetaItem(audio=audio_buffer, source_language=src_lang)
        self._whisper_pipe.input_queue.put(item)
        return self._whisper_pipe.output_queue.get()
    
    def voice_detect(self, audio_buffer:bytes) -> MetaItem:
        item = MetaItem(source_audio=audio_buffer)
        self._vad_pipe.input_queue.put(item)
        return self._vad_pipe.output_queue.get()

    

if __name__ == "__main__":
    import soundfile
    tp = TranslatePipes()
    # result = tp.translate("你好，今天天气怎么样?", src_lang="zh", dst_lang="en")
    mel, _, = soundfile.read("assets/jfk.flac")
    # result = tp.transcrible(mel, 'en')
    result = tp.voice_detect(mel)
    print(result)