import gradio as gr
from transformers import pipeline
from models.jairvis import JAIRVISMKV
from utils.audio_utils import text_to_speech, speech_to_text
from utils.memory_manager import MemoryManager
from config import MODEL_NAME, API_KEY

class JairvisApp:
    def __init__(self):
        self.jairvis = JAIRVISMKV()
        self.memory_manager = MemoryManager()
        self.nlp_pipeline = pipeline("text-generation", model=MODEL_NAME)

    def process_input(self, input_text):
        # Process input and generate response
        response = self.jairvis.generate_response(input_text)
        
        # Update memory
        self.memory_manager.store(input_text, response)
        
        # Convert response to speech
        audio_output = text_to_speech(response)
        
        return response, audio_output

    def voice_input(self, audio):
        # Convert speech to text
        input_text = speech_to_text(audio)
        
        # Process the text input
        response, audio_output = self.process_input(input_text)
        
        return input_text, response, audio_output

# Set up the Gradio interface
def launch():
    app = JairvisApp()
    
    text_interface = gr.Interface(
        fn=app.process_input,
        inputs="text",
        outputs=["text", "audio"],
        title="JAIRVISMKV Text Interface"
    )
    
    voice_interface = gr.Interface(
        fn=app.voice_input,
        inputs="audio",
        outputs=["text", "text", "audio"],
        title="JAIRVISMKV Voice Interface"
    )
    
    combined_interface = gr.TabbedInterface([text_interface, voice_interface], ["Text", "Voice"])
    combined_interface.launch(share=True)

if __name__ == "__main__":
    launch()