import gradio as gr
import openai
import os

# Retrieve the API key from the Hugging Face secret
api_key = os.environ["INSTRUCTOR_API_KEY"]
instructor_prompt = os.environ["INSTRUCTOR_PROMPT"]

with gr.Blocks() as demo:
    chat_history = gr.State(value=[])

    def chat_with_instructor(message, history):

        openai.api_key = api_key
        chat_input = instructor_prompt + "\nUser: " + message + "\nInstructor:"
        response = openai.Completion.create(
            engine="text-davinci-002",
            prompt=chat_input,
            max_tokens=50  # Adjust this for response length
        ).choices[0].text.strip()

        chat_history.value.append(
            {
                "user": message,
                "instructor": instructor_prompt,
                "bot": response,
            }
        )
        return response

    def generate_json(chat_history):
        return chat_history.value

    gr.ChatInterface(
        fn=chat_with_instructor,

        title="Chat with Instructor",
        description="Chat with the instructor using a predefined prompt.",
        examples=["Hi, can you help me with this?"],

        submit_btn="Send",
        stop_btn="Stop",
        retry_btn="Retry",
        undo_btn="Undo last message",
        clear_btn="Start a new conversation"
    ).queue()

    chat_history_json = gr.JSON(generate_json(chat_history))
    gr.Markdown("### 📩 Generate the JSON file for your chat history!")
    gr.Interface(fn=generate_json,
                 inputs=None,
                 outputs=[chat_history_json])

# demo.queue()
demo.launch()