import gradio as gr
import numpy as np
import random
import time
import os
import shutil
import codecs
# How to RUN code ==> gradio gradio_llm_example.py 


# Define text and title information 
title1 = "##  </br> </br> </br> 🤗💬 QA App"

title2 = " ##  </br> </br> </br> Gradio QA Bot"

intro = """ Welcome! This is not just any bot, it's a special one equipped with state-of-the-art natural language processing capabilities, and ready to answer your queries.


    Ready to explore? Let's get started!
  
    
    * Step 1: Upload a PDF document.
    * Step 2: Type in a question related to your document's content.
    * Step 3: Get your answer!
 
    
    Push clear cache before uploading a new doc!

    
    """

about = """ 
    ## </br> About
    This app is an LLM-powered chatbot built using:
    - [Streamlit](<https://streamlit.io/>)
    - [HugChat](<https://github.com/Soulter/hugging-chat-api>)
    - Chat Model = llama2-chat-hf 7B 
    - Retreiver model = all-MiniLM-L6-v2

    </br> 
    💡 Note: No API key required!
    
     </br> 
    Made with ❤️ by us
     """


# Define theme ==> see gr.themes.builder()
theme = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="emerald",
    neutral_hue="slate",
).set(
    body_background_fill_dark='*primary_50',
    shadow_drop='*shadow_spread',
    button_border_width='*block_border_width',
    button_border_width_dark='*block_label_border_width'
)


def upload_file(files_obj):
    """ Upload several files from drag and drop, and save them in local temp folder
    files_obj (type:list) : list of tempfile._TemporaryFileWrapper 
    return checkbox to display uploaded documents """
    # Create local copy
    temp_file_path = "./temp"
    if not os.path.exists(temp_file_path):
        os.makedirs(temp_file_path)
    # Save each file among list of given files
    file_name_list = list()
    for file_obj in files_obj :
        file_name = os.path.basename(file_obj.name)
        file_name_list.append(file_name)
        shutil.copyfile(file_obj.name, os.path.join(temp_file_path, file_name))
    return {uploaded_check : gr.Radio(choices=file_name_list, visible=True),
            choose_btn : gr.Button(value="Choose", visible=True)}


def read_content(file_name):
    print(file_name, type(file_name))
    temp_file_path = "./temp"
    file_path = os.path.join(temp_file_path, file_name)
    with open(file_path, "rb") as file:
        try:
            content = file.read()
            print(content)
            print(codecs.decode(content, 'utf-8'))
            return {error_box: gr.Textbox(value=f"File ready to be used. \n You can ask a question about the uploaded PDF document.", visible=True)}
        except Exception as e:
            print(f"Error occurred while writing the file: {e}")
            return {error_box: gr.Textbox(value=f"Error occurred while writing the file: {e}", visible=True)}


def respond(message, chat_history, 
            language_choice, max_length, temperature,
            num_return_sequences, top_p, no_repeat_ngram_size):
        #No LLM here, just respond with a random pre-made message
        if content == "":
            bot_message = f"j'ai {max_length}" + random.choice(["Tell me more about it", 
                                            "Cool, but I'm not interested", 
                                            "Hmmmm, ok then"]) 
        else:
            bot_message = " j'ai besoin d'un modèle pour lire le {content[:3]}"
        chat_history.append((message, bot_message))
        return "", chat_history



# Layout
with gr.Blocks(theme=gr.themes.Soft()) as gradioApp:
    with gr.Row():
        with gr.Column(scale=1, min_width=100):
            logo_gr = gr.Markdown(""" <img src="file/logo_neovision.png" alt="logo" style="width:400px;"/>""") 
            # gr.Image("./logo_neovision.png")
            about_gr = gr.Markdown(about)   
            
        with gr.Column(scale=2, min_width=500):
            title1_gr= gr.Markdown(title1)
            intro_gr = gr.Markdown(intro)

            #  Upload several documents
            content = ""
            upload_button = gr.UploadButton("Browse files", label="Drag and drop your documents here", 
                                            size="lg", scale=0, min_width=100,
                                            file_types=["pdf"], file_count="multiple")
            uploaded_check = gr.Radio(label="Uploaded documents", visible=False,
                                      info="Do you want to use a supporting document?")
            choose_btn = gr.Button(value="Choose", visible=False)
            upload_button.upload(upload_file, upload_button, [uploaded_check, choose_btn])

            # Read only one document
            error_box = gr.Textbox(label="Reading files... ", visible=False)
            choose_btn.click(read_content, inputs=uploaded_check, outputs=error_box)

            # Select advanced options
            gr.Markdown(""" ## Toolbox """)
            with gr.Accordion(label="Select advanced options",open=False):
                language_choice = gr.Dropdown(["English", "French"], label="Language", info="Choose your language")
                max_length = gr.Slider(label="Token length", minimum=1, maximum=100, value=50, step=1)
                temperature= gr.Slider(label="Temperature", minimum=0.1, maximum=1, value=0.8, step=0.1)
                num_return_sequences= gr.Slider(label="Temperature", minimum=0.1, maximum=50, value=1, step=0.1)
                top_p= gr.Slider(label="Temperature", minimum=0.1, maximum=1, value=0.8, step=0.1)
                no_repeat_ngram_size= gr.Slider(label="Temperature", minimum=0.1, maximum=1, value=3, step=0.1)


        # Chat
        with gr.Column(scale=2, min_width=600):
            title2_gr = gr.Markdown(title2)
            chatbot = gr.Chatbot(label="Bot", height=500) 
            msg = gr.Textbox(label="User", placeholder="Ask any question.")
            chatbot_btn = gr.Button("Submit")
            chatbot_btn.click(respond, inputs=[msg, chatbot, 
                                               language_choice, max_length, temperature,
                                               num_return_sequences, top_p, no_repeat_ngram_size],
                                               outputs=[msg, chatbot])
            clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")


gr.close_all()
gradioApp.launch(share=True, enable_queue=True)