import os
import gradio as gr
import openai
import pandas as pd

# The OpenAI API key provided by the instructor
openai.api_key = os.environ.get("OPENAI_API_KEY")
# Instructor's prompt (invisible to students)
instructor_prompt = os.environ.get("SECRET_PROMPT")

def add_user_message(msg, history, messages):
    messages.append({"role": "user", "content": msg})
    return "", history + [[msg, None]], messages
    
def get_tutor_reply(history, messages):
    # Generate a response from the chatbot
    completion = openai.chat.completions.create( 
        model="gpt-3.5-turbo",
        messages=messages,
        max_tokens=150
    ) 
    # Extract the chatbot's reply
    reply = completion.choices[0].message.content 
    history[-1][1] = reply
    messages.append({"role": "assistant", "content": reply})
    return history, messages

def save_chat_to_json(history):
    formatted_convo = pd.DataFrame(history, columns=['user', 'chatbot'])
    output_fname = f'tutoring_conversation.json'
    formatted_convo.to_json(output_fname, orient='records')
    return gr.update(value=output_fname, visible=True)
    
# Create a Gradio interface
with gr.Blocks() as SimpleChat:
    gr.Markdown("""
        ## Chat with A Tutor
        Description here
        """)
    messages = gr.JSON(
        visible=False,
        value = [{"role": "system", "content":  instructor_prompt}]
    )
    with gr.Group():
        chatbot = gr.Chatbot(label="Simple Tutor")
        with gr.Row(equal_height=True):
            user_chat_input = gr.Textbox(show_label=False, scale=9)
            submit_btn = gr.Button("Enter", scale=1)
    with gr.Group():
        export_dialogue_button_json = gr.Button("Export your chat history as a .json file")
        file_download = gr.Files(label="Download here", file_types=['.json'], visible=False)
    
    submit_btn.click(
        fn=add_user_message, inputs=[user_chat_input, chatbot, messages], outputs=[user_chat_input, chatbot, messages], queue=False
    ).then(
        fn=get_tutor_reply, inputs=[chatbot, messages], outputs=[chatbot, messages], queue=True
    )
    user_chat_input.submit(
        fn=add_user_message, inputs=[user_chat_input, chatbot, messages], outputs=[user_chat_input, chatbot, messages], queue=False
    ).then(
        fn=get_tutor_reply, inputs=[chatbot, messages], outputs=[chatbot, messages], queue=True
    )
    export_dialogue_button_json.click(save_chat_to_json, chatbot, file_download, show_progress=True)

# Launch the Gradio app
if __name__ == "__main__":
    SimpleChat.launch()