import gradio as gr
import pandas as pd
from functools import partial

def save_chatbot_dialogue(chat_tutor, save_type):

    formatted_convo = pd.DataFrame(chat_tutor.conversation_memory, columns=['user', 'chatbot'])

    output_fname = f'tutoring_conversation.{save_type}'

    if save_type == 'csv':
        formatted_convo.to_csv(output_fname, index=False)
    elif save_type == 'json':
        formatted_convo.to_json(output_fname, orient='records')
    elif save_type == 'txt':
        temp = formatted_convo.apply(lambda x: 'User: {0}\nAI: {1}'.format(x[0], x[1]), axis=1)
        temp = '\n\n'.join(temp.tolist())
        with open(output_fname, 'w') as f:
            f.write(temp)
    else:
        gr.update(value=None, visible=False)
    
    return gr.update(value=output_fname, visible=True)

save_json = partial(save_chatbot_dialogue, save_type='json')
save_txt = partial(save_chatbot_dialogue, save_type='txt')


# history is a list of list  
# [[user_input_str, bot_response_str], ...]

class BasicTutor:
    # create basic initialization function
    def __init__(self):
        self.conversation_memory = []
        self.flattened_conversation = ''

    def add_user_message(self, user_message):
        self.conversation_memory.append([user_message, None])
        self.flattened_conversation = self.flattened_conversation + '\n\n' + 'User: ' + user_message

    def get_tutor_reply(self, user_message):
        # get tutor message
        tutor_message = "You said: " + user_message
        # add tutor message to conversation memory
        self.conversation_memory[-1][1] = tutor_message
        self.flattened_conversation = self.flattened_conversation + '\nAI: ' + tutor_message

    def forget_conversation(self):
        self.conversation_memory = []
        self.flattened_conversation = ''


### Chatbot Functions ###
def add_user_message(user_message, chat_tutor):
  """Display user message and update chat history to include it."""
  chat_tutor.add_user_message(user_message)
  return chat_tutor.conversation_memory, chat_tutor

def get_tutor_reply(user_message, chat_tutor):
  chat_tutor.get_tutor_reply(user_message)
  return gr.update(value="", interactive=True), chat_tutor.conversation_memory, chat_tutor


with gr.Blocks() as demo:
    #initialize tutor (with state)
    study_tutor = gr.State(BasicTutor())
        
    # Chatbot interface
    gr.Markdown("""
    ## Chat with the Model
    Description here
    """)

    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)
    
    user_chat_submit.click(add_user_message, 
                           [user_chat_input, study_tutor], 
                           [chatbot, study_tutor], queue=False).then(
        get_tutor_reply, [user_chat_input, study_tutor], [user_chat_input, chatbot, study_tutor], queue=True)
    
    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")
        
        file_download = gr.Files(label="Download here",
                                file_types=['.txt', '.json'], 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)
    
demo.queue()
demo.launch()