import gradio as gr
import os

# Install the groq package if it is not installed
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq

# Set up the Groq client with the secret key
groq_key = os.getenv('groq_key')
if not groq_key:
    raise ValueError("groq_key environment variable is not set")

client = Groq(api_key=groq_key)

class SimpleChatBot:
    def __init__(self):
        self.initial_prompt = [
            {
                "role": "system",
                "content": "你是一個英文老師，會教我背單字，並用蘇格拉底問答引導我，引導時請用繁體中文 zhTW"
            }
        ]

    def get_response(self, message, chat_history):
        messages = self.initial_prompt + chat_history
        messages.append({"role": "user", "content": message})
        
        completion = client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=messages,
            temperature=1,
            max_tokens=1024,
            top_p=1,
            stream=True,
            stop=None,
        )
        
        response_content = ""
        for chunk in completion:
            response_content += chunk.choices[0].delta.content or ""
        
        return response_content

chatbot = SimpleChatBot()

def respond(message, chat_history):
    chat_history = [{"role": entry["role"], "content": entry["content"]} for entry in chat_history]
    response = chatbot.get_response(message, chat_history)
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": response})
    return chat_history, ""

with gr.Blocks(title="簡單的Gradio聊天機器人") as demo:
    gr.Markdown("# 簡單的Gradio聊天機器人")
    
    chatbot_interface = gr.Chatbot(type="messages")
    
    with gr.Row():
        user_input = gr.Textbox(placeholder="輸入訊息...", label="你的訊息")
        send_button = gr.Button("發送")
    
    send_button.click(respond, inputs=[user_input, chatbot_interface], outputs=[chatbot_interface, user_input])
    
    demo.launch(share=True)