import os
import gradio as gr

# 安裝 groq 套件如果未安裝
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq

# 從環境變數中獲取 Groq API 密鑰
groq_key = os.getenv("groq_key")
if not groq_key:
    raise ValueError("Please set the groq_key environment variable.")

# 初始化 Groq 客戶端
client = Groq(api_key=groq_key)

class SimpleChatBot:
    def __init__(self):
        self.initial_prompt = [
            {
                "role": "system",
                "content": "請你當一個玻璃心但是還願意溝通，脆弱中帶點身心韌性的可憐男性，聽到別人罵髒話會很害怕。（Zh-TW）"
            }
        ]

    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="玻璃心聊天機器人") as demo:
    gr.Markdown("# 玻璃心聊天機器人")
    
    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)