import gradio as gr
import edge_tts
import asyncio
import tempfile
import os
from huggingface_hub import InferenceClient
import re
from streaming_stt_nemo import Model
import torch
import random

# Default language for speech-to-text
default_lang = "en"

# Initialize STT engine
engines = { default_lang: Model(default_lang) }

def transcribe(audio):
    """
    Transcribes audio file to text.
    """
    lang = "en"
    model = engines[lang]
    text = model.stt_file(audio)[0]
    return text

# Get Hugging Face token from environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)

def client_fn(model):
    """
    Selects the appropriate model based on user choice.
    """
    if "Mixtral" in model:
        return InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
    elif "Llama" in model:
        return InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")
    elif "Mistral" in model:
        return InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
    elif "Phi" in model:
        return InferenceClient("microsoft/Phi-3-mini-4k-instruct")
    else: 
        return InferenceClient("microsoft/Phi-3-mini-4k-instruct")

def randomize_seed_fn(seed: int) -> int:
    """
    Randomizes the seed for generating responses.
    """
    seed = random.randint(0, 999999)
    return seed

# System instructions for the virtual assistant
system_instructions = ("[SYSTEM] You are a virtual assistant named Julia, made by 'Harshana Lakshara.'"
                       "Your task is to Answer the question."
                       "Keep conversation very short, clear and concise."
                       "Respond naturally and concisely to the user's queries. "
                       "Respond in a normal, conversational manner while being friendly and helpful."
                       "The expectation is that you will avoid introductions and start answering the query directly, Only answer the question asked by user, Do not say unnecessary things."
                       "Begin with a greeting if the user initiates the conversation. "
                       "Avoid unnecessary introductions and answer the user's questions directly."
                       "Here is the user's query:[QUESTION] ")

def models(text, model="Mixtral 8x7B", seed=42):
    """
    Generates a response based on user input and model selection.
    """
    seed = int(randomize_seed_fn(seed))
    generator = torch.Generator().manual_seed(seed)  
    
    client = client_fn(model)
    
    generate_kwargs = dict(
        max_new_tokens=300,
        seed=seed
    )    
    formatted_prompt = system_instructions + text + "[JARVIS]"
    stream = client.text_generation(
        formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        if not response.token.text == "</s>":
            output += response.token.text
    return output

async def respond(audio, text, model, seed):
    """
    Handles user input (audio or text), generates a response, 
    and converts the response to audio.
    """
    if audio:
        user = transcribe(audio)
    else:
        user = text
    reply = models(user, model, seed)
    communicate = edge_tts.Communicate(reply)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
        tmp_path = tmp_file.name
        await communicate.save(tmp_path)
    return reply, tmp_path

# Description of the Gradio interface
DESCRIPTION = """ # <center><b>JARVIS⚡</b></center>
        ### <center>A personal Assistant of Tony Stark for YOU
        ### <center>Voice Chat with your personal Assistant</center>
        """

# Gradio interface
with gr.Blocks(css="style.css") as demo:    
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        select = gr.Dropdown([ 'Mixtral 8x7B',
        'Llama 3 8B',
        'Mistral 7B v0.3',
        'Phi 3 mini',
    ],
    value="Mistral 7B v0.3",
    label="Model"
    )
        seed = gr.Slider(
        label="Seed",
        minimum=0,
        maximum=999999,
        step=1,
        value=0,
        visible=False
        )
        input_audio = gr.Audio(label="User (Audio)", sources="microphone", type="filepath", waveform_options=False)
        input_text = gr.Textbox(label="User (Text)", placeholder="Type your message here...")
        send_button = gr.Button("Send")
        output_text = gr.Textbox(label="AI (Text)", interactive=False)
        output_audio = gr.Audio(label="AI (Audio)", type="filepath",
                        interactive=False,
                        autoplay=True,
                        elem_classes="audio")
        
        # Set up button click event
        send_button.click(
            fn=respond, 
            inputs=[input_audio, input_text, select, seed],
            outputs=[output_text, output_audio]
        )

if __name__ == "__main__":
    demo.queue(max_size=200).launch()