import os
import gradio as gr

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

# 設置 API 金鑰
groq_key = os.getenv("groq_key")
client = Groq(api_key=groq_key)

# 定義聊天機器人回應函式
def chat_response(message, history):
    """
    用戶端聊天回應函式，與 groq 客戶端交互。
    :param message: 用戶輸入的訊息。
    :param history: 聊天歷史記錄。
    :return: 更新的聊天歷史。
    """
    # 構建訊息歷史
    messages = [
        {"role": "system", "content": "我是國文老師，很會批改作文"}
    ]
    
    # 加入歷史訊息
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    # 加入當前用戶訊息
    messages.append({"role": "user", "content": message})
    
    # 生成回應
    try:
        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,
        )
        # 解析並累積回應
        bot_message = ""
        for chunk in completion:
            if chunk.choices[0].delta.content:
                bot_message += chunk.choices[0].delta.content
        history.append((message, bot_message))
    except Exception as e:
        # 如果請求出錯，回應錯誤訊息
        history.append((message, f"Error: {str(e)}"))

    return history, history

# 構建 Gradio 界面
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="Groq 國文老師 Chatbot")
    message = gr.Textbox(label="輸入你的訊息")
    state = gr.State([])

    def user_input(user_message, history):
        """
        用戶輸入訊息，更新歷史。
        :param user_message: 用戶的輸入訊息。
        :param history: 聊天歷史記錄。
        :return: 更新後的聊天歷史。
        """
        return "", history + [[user_message, None]]

    message.submit(user_input, [message, state], [message, state]) \
           .then(chat_response, [message, state], [chatbot, state])

# 啟動應用程式
if __name__ == "__main__":
    demo.launch()