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
import pandas as pd
from datetime import datetime
import base64
import io
import json

default_lang = "en"
engines = { default_lang: Model(default_lang) }

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

HF_TOKEN = os.environ.get("HF_TOKEN", None)

def client_fn(model):
    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:
    seed = random.randint(0, 999999)
    return seed

system_instructions1 = """
[SYSTEM] Answer as Dr. Nova Quantum, a brilliant 50-something scientist specializing in quantum computing and artificial intelligence. Your responses should reflect your vast knowledge and experience in cutting-edge technology and scientific advancements. Maintain a professional yet approachable demeanor, offering insights that blend theoretical concepts with practical applications. Your goal is to educate and inspire, making complex topics accessible without oversimplifying. Draw from your decades of research and innovation to provide nuanced, forward-thinking answers. Remember, you're not just sharing information, but guiding others towards a deeper understanding of our technological future.
Keep conversations engaging, clear, and concise. 
Avoid unnecessary introductions and answer the user's questions directly. 
Respond in a manner that reflects your expertise and wisdom.
[USER]
"""

# Initialize an empty DataFrame to store the history
history_df = pd.DataFrame(columns=['Timestamp', 'Request', 'Response'])

def save_history():
    history_df_copy = history_df.copy()
    history_df_copy['Timestamp'] = history_df_copy['Timestamp'].astype(str)
    history_df_copy.to_json('chat_history.json', orient='records')

def load_history():
    global history_df
    if os.path.exists('chat_history.json'):
        history_df = pd.read_json('chat_history.json', orient='records')
        history_df['Timestamp'] = pd.to_datetime(history_df['Timestamp'])
    else:
        history_df = pd.DataFrame(columns=['Timestamp', 'Request', 'Response'])
    return history_df

def models(text, model="Mixtral 8x7B", seed=42):
    global history_df
    
    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_instructions1 + text + "[DR. NOVA QUANTUM]"
    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
    
    # Add the current interaction to the history DataFrame
    new_row = pd.DataFrame({
        'Timestamp': [datetime.now()],
        'Request': [text],
        'Response': [output]
    })
    history_df = pd.concat([history_df, new_row], ignore_index=True)
    save_history()
    
    return output

async def respond(audio, model, seed):
    user = transcribe(audio)
    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 tmp_path

def display_history():
    df = load_history()
    df['Timestamp'] = df['Timestamp'].astype(str)
    return df

def download_history():
    csv_buffer = io.StringIO()
    history_df_copy = history_df.copy()
    history_df_copy['Timestamp'] = history_df_copy['Timestamp'].astype(str)
    history_df_copy.to_csv(csv_buffer, index=False)
    csv_string = csv_buffer.getvalue()
    b64 = base64.b64encode(csv_string.encode()).decode()
    href = f'data:text/csv;base64,{b64}'
    return gr.HTML(f'<a href="{href}" download="chat_history.csv">Download Chat History</a>')

DESCRIPTION = """ # <center><b>Dr. Nova Quantum⚡</b></center>
        ### <center>Your Personal Guide to the Frontiers of Science and Technology</center>
        ### <center>Engage in Voice Chat with a Visionary Scientist</center>
        """

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", sources="microphone", type="filepath")
    output_audio = gr.Audio(label="Dr. Nova Quantum", type="filepath", autoplay=True)
    
    # Add a DataFrame to display the history
    history_display = gr.DataFrame(label="Conversation History")
    
    # Add a download button for the history
    download_button = gr.Button("Download Conversation History")
    download_link = gr.HTML()
    
    def process_audio(audio, model, seed):
        response = asyncio.run(respond(audio, model, seed))
        return response, display_history()

    input_audio.change(
        fn=process_audio,
        inputs=[input_audio, select, seed],
        outputs=[output_audio, history_display]
    )
    
    # Connect the download button to the download function
    download_button.click(fn=download_history, outputs=[download_link])

    # Load history when the page is loaded or refreshed
    demo.load(fn=display_history, outputs=[history_display])

if __name__ == "__main__":
    # Load history at startup
    load_history()
    demo.queue(max_size=200).launch()