import os

# 嘗試安裝並導入 groq
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')  # 嘗試安裝 groq
    from groq import Groq

import gradio as gr

# 從環境變數中取得 API key
groq_key = os.getenv('groq_key')

# 創建 Groq 客戶端
client = Groq(api_key=groq_key)

# 定義與 Groq 的聊天功能
def chat_with_groq(message, chat_history):
    # 將先前的對話記錄加到 messages 中
    messages = [{"role": "system", "content": "你是一位國小老師，使用的語言為繁體中文(zh-tw)。你的專長是將生活素養的正確觀念融入課堂教學中，無論學生問你什麼問題，你都會把話題引導到正向觀念的討論上。"}]
    for user_msg, bot_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": message})

    # 呼叫 groq API 生成回應
    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",  # 使用 groq 的模型
        messages=messages,
        temperature=1,
        max_tokens=3660,
        top_p=1,
        stream=False,
        stop=None,
    )

    # 獲取模型回應
    response = completion['choices'][0]['message']['content']
    
    # 更新聊天記錄
    chat_history.append((message, response))
    
    return chat_history, chat_history

# 建立 Gradio 的聊天機器人介面
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()  # 聊天機器人介面
    user_input = gr.Textbox(placeholder="輸入你的問題", label="輸入")
    send_button = gr.Button("發送")  # 發送按鈕
    clear_button = gr.Button("清除對話")  # 清除聊天紀錄按鈕
    chat_history = gr.State([])  # 儲存聊天紀錄的狀態

    # 設定按鈕點擊後的動作：呼叫 `chat_with_groq`，將結果顯示在 `chatbot` 中
    send_button.click(chat_with_groq, inputs=[user_input, chat_history], outputs=[chatbot, chat_history])
    clear_button.click(lambda: None, None, chatbot, queue=False)  # 清除聊天內容

# 啟動 Gradio 介面
demo.launch()