import time
import yaml
import gradio as gr
import os
import argparse
import random
from chatbot import ChatBot

# MARKDOWN
MARKDOWN = """
# Coffee-Gym Feedback Model
Welcome to the COFFEE-GYM demo page! This page will guide you through using our comprehensive RL environment for training models that provide feedback on code editing.

## Prompt template
To use the COFFEE-GYM feedback model, you can follow the prompt template below:

~~~json
Problem Description:
{problem}

Incorrect Code:
```python
{code}
```

Please generate feedback for the wrong code.
~~~

## Response
The chatbot will provide feedback on the incorrect code by analyzing the problem description, input, output, and the buggy solution provided in the prompt. You can interact with the chatbot to get immediate feedback, regenerate responses, and clear the chat history.

Feel free to explore and make the most out of COFFEE-GYM!
"""


def main(args):
    
    chatbot = ChatBot()
    
    with gr.Blocks() as app:
        ##### Playground #####
        gr.Markdown(MARKDOWN)
        
        chat_history = gr.Chatbot(show_copy_button=True)     
        input_text = gr.Textbox(label="Enter your prompt and click 'Send'", placeholder="Type your message...")
        send_button = gr.Button("Send", elem_id="send-btn")
        
        with gr.Row():
            regenerate_button = gr.Button("Regenerate")
            clear_button = gr.Button("Clear")

        send_button.click(
            fn=chatbot.chat, 
            inputs=[chat_history, input_text], 
            outputs=chat_history
        )
        regenerate_button.click(
            fn=chatbot.regenerate, 
            inputs=[chat_history, input_text], 
            outputs=chat_history
        )

        clear_button.click(
            fn=chatbot.clear_chat, 
            inputs=[], 
            outputs=chat_history
        )
            
    app.launch(share=True)


if __name__ == "__main__":
    main()