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

# MARKDOWN
MARKDOWN = """
# Coffee-Gym
Welcome to the COFFEE-GYM demo page! 

## DS-Coder-7B-PPO-CoffeeEval
You can try using the COFFEE-GYM feedback model (DS-Coder-7B-PPO-CoffeeEval) to get feedback on your code.

## Prompt template
To use the DS-Coder-7B-PPO-CoffeeEvall, 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!
"""




# get arguments from command line
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, default="Team-Coffee-Gym/DS-Coder-7B-PPO-CoffeeEval", help="Model name to use for the chatbot.")
    parser.add_argument("--url", type=str, default="https://critic.jp.ngrok.io/v1", help="URL to use for the chatbot.")
    args = parser.parse_args()  
    return args


def main(args):
    
    chatbot = ChatBot(model_name=args.model_name, url=args.url)
    
    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__":
    args = get_args()
    main(args)