import os
import gradio as gr
import pandas as pd
from functools import partial
from ai_classroom_suite.MediaVectorStores import *
from ai_classroom_suite.UIBaseComponents import *

#  default folder path
folder_path = "context_files"
#  default output file name
out_file_name = "vector_store.txt"

# Check if vector store file already exist on disk
def vector_store_file_exist():
    # Get all files in the folder
    files = os.listdir(folder_path)
    # Check if output file already exist in this folder
    return (out_file_name in files)

# Helper function to get all files' paths from a folder
# Return a list of file paths except for README.txt and vector_store.txt (if exist)
def get_filepaths_from_folder(folder_path):
    # Store the paths of files
    filepath_list = []
    
    # Check if the specified folder exists
    if not os.path.exists(folder_path):
        print(f"Folder '{folder_path}' does not exist.")
        return filepath_list
    
    # Get all the files in the folder
    files = os.listdir(folder_path)
    
    for file_name in files:
        # Excluding README.txt and vector_store.txt
        if file_name != "README.txt" and file_name != "vector_store.txt":
            # Get the file path for each item
            file_path = os.path.join(folder_path, file_name)
            # Check if the item is a file and not a subdirectory
            if os.path.isfile(file_path):
                filepath_list.append(file_path)
    
    return filepath_list

# Helper function to write content of files in a folder to output file
def write_vector_store_to_file(out_file_name):
    # If vector_store.txt already exist, return nothing
    if vector_store_file_exist():
        return gr.File(value=out_file_name, visible=False)
 
    # Only try to create the vector store if vector_store.txt doesn't exist
    else:
        # Call the function to read files (excluding README.txt and vector_store.txt) pathes
        filepath_list = get_filepaths_from_folder(folder_path)
        # Extract the text out from files
        files_content = files_to_text(filepath_list, chunk_size=100, chunk_overlap=20)
        
        # Write the vector_store onto the output file
        with open(out_file_name, "w") as f:
            for i in range(len(files_content)):
                item = str(files_content[i]) + "\n"
                f.write(item)
        
        # Show the downlodable vector store file and give instruction on upload the vector store file to disk (on HuggingFace)
        return gr.File(title="Download your vector store file and upload it into the context_files folder under Files", 
                       value=out_file_name, visible=True)

# overwrites the original method since we don't deal with any vector stores display here
def get_tutor_reply(chat_tutor):
    chat_tutor.get_tutor_reply()
    return gr.update(value="", interactive=True), chat_tutor.conversation_memory, chat_tutor

def get_conversation_history(chat_tutor):
    return chat_tutor.conversation_memory, chat_tutor

# To show the loading process on the button when creating vector store file
def creating_vs_button(obj_in):
    return gr.update(interactive=False, value='Creating Vector Store file...')
# To show the loading process on the button when initializing tutor
def initializing_tutor_button(obj_in):
    return gr.update(interactive=False, value='Initializing Tutor...')


with gr.Blocks() as ReadingQuiz:
    #initialize tutor (with state)
    study_tutor = gr.State(SlightlyDelusionalTutor())

    # Student chatbot interface
    gr.Markdown("""
    ## Chat with the Model
    This is the Blocher Reading Quiz App.
    """)
    
    # Instead of ask students to provide key, the key is now provided by the instructor. 
    api_input = gr.Textbox(show_label=False, type="password", visible=False, value=os.environ.get("OPENAI_API_KEY"))

    # The instructor will provide a secret prompt/persona to the tutor
    instructor_prompt = gr.Textbox(label="Verify your prompt content", value = os.environ.get("SECRET_PROMPT"), visible=False)
    
    # Show input files
    file_input = gr.File(label="Reading materials", value=get_filepaths_from_folder(folder_path), visible=True)
    
    # Show output file for vector store when needed
    vs_file_name = gr.Text(visible=False, value=out_file_name)
    file_output = gr.File(visible=False)
    
    # Placeholders components
    text_input_none = gr.Textbox(visible=False)
    file_input_none = gr.File(visible=False)
    instructor_input_none = gr.TextArea(visible=False)
    learning_objectives_none = gr.Textbox(visible=False)

    # Set the secret prompt in this session and embed it to the study tutor
    vs_build_button = gr.Button("Initialize Tutor")
    vs_build_button.click(
        fn=creating_vs_button, inputs=vs_build_button, outputs=vs_build_button
    ).then(
        fn=write_vector_store_to_file, inputs=[vs_file_name], outputs=[file_output]
    ).then(
        fn=initializing_tutor_button, inputs=[vs_build_button], outputs=[vs_build_button]
    ).then(
        fn=create_reference_store, 
        inputs=[study_tutor, vs_build_button, instructor_prompt, file_output, instructor_input_none, api_input, learning_objectives_none],
        outputs=[study_tutor, vs_build_button]
    )

    with gr.Row(equal_height=True):
        with gr.Column(scale=2):
            chatbot = gr.Chatbot()
            with gr.Row():
                user_chat_input = gr.Textbox(label="User input", scale=9)
                user_chat_submit = gr.Button("Ask/answer model", scale=1)

    # First add user's message to the conversation history
    # Then get reply from the tutor and add that to the conversation history
    user_chat_submit.click(
        fn = add_user_message, inputs = [user_chat_input, study_tutor], outputs = [user_chat_input, chatbot, study_tutor], queue=False
    ).then(
        fn = get_tutor_reply, inputs = [study_tutor], outputs = [user_chat_input, chatbot, study_tutor], queue=True
    )
    # User can also press "Enter" on keyboard to submit a message
    user_chat_input.submit(
        fn = add_user_message, inputs = [user_chat_input, study_tutor], outputs = [user_chat_input, chatbot, study_tutor], queue=False
    ).then(
        fn = get_tutor_reply, inputs = [study_tutor], outputs = [user_chat_input, chatbot, study_tutor], queue=True
    )

    # Download conversation history file
    with gr.Blocks():
        gr.Markdown("""
        ## Export Your Chat History
        Export your chat history as a .json, .txt, or .csv file
        """)
        with gr.Row():
            export_dialogue_button_json = gr.Button("JSON")
            export_dialogue_button_txt = gr.Button("TXT")
            export_dialogue_button_csv = gr.Button("CSV")
            
        file_download = gr.Files(label="Download here", file_types=['.json', '.txt', '.csv'], type="file", visible=False)
    
    export_dialogue_button_json.click(save_json, study_tutor, file_download, show_progress=True)
    export_dialogue_button_txt.click(save_txt, study_tutor, file_download, show_progress=True)
    export_dialogue_button_csv.click(save_csv, study_tutor, file_download, show_progress=True)
    
ReadingQuiz.queue().launch(server_name='0.0.0.0', server_port=7860)