import gradio as gr
import os
import json
import requests
from typing import List, Dict

# 環境變數設置
XAI_API_KEY = os.getenv("XAI_API_KEY")
if not XAI_API_KEY:
    raise ValueError("Please set XAI_API_KEY environment variable")

def call_grok_api(messages: List[Dict]) -> str:
    """調用Grok API"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {XAI_API_KEY}"
    }
    
    data = {
        "messages": messages,
        "model": "grok-beta",
        "stream": False,
        "temperature": 0.7  # 可以調整此值來控制回答的創造性
    }
    
    try:
        response = requests.post(
            "https://api.x.ai/v1/chat/completions",
            headers=headers,
            json=data
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]
    except requests.exceptions.RequestException as e:
        return f"Error calling Grok API: {str(e)}"

def format_history(history: List[Dict]) -> List[Dict]:
    """將Gradio的歷史記錄格式化為Grok API需要的格式"""
    formatted_messages = [
        {
            "role": "system",
            "content": "You are Grok, a chatbot inspired by the Hitchhikers Guide to the Galaxy."
        }
    ]
    
    for msg in history:
        formatted_messages.append({
            "role": "user" if msg["role"] == "user" else "assistant",
            "content": msg["content"]
        })
    
    return formatted_messages

def respond(message: str, history: List[Dict]) -> Dict:
    """處理用戶輸入並生成回應"""
    history.append({"role": "user", "content": message})
    
    # 格式化消息歷史
    formatted_history = format_history(history)
    
    # 獲取Grok的回應
    response = call_grok_api(formatted_history)
    
    # 添加助手回應到歷史記錄
    history.append({"role": "assistant", "content": response})
    
    return history

# 創建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("""
    # Grok Chatbot
    與Elon Musk的Grok AI聊天！這個聊天機器人使用了xAI的Grok API。
    """)
    
    chatbot = gr.Chatbot(
        value=[],
        bubble_full_width=False,
        avatar_images=(None, "🤖"),
        height=500
    )
    
    msg = gr.Textbox(
        placeholder="輸入訊息...",
        container=False,
        scale=7
    )
    
    clear = gr.ClearButton([msg, chatbot], variant="secondary", scale=1)
    
    msg.submit(respond, [msg, chatbot], [chatbot], api_name="chat")
    
demo.queue()

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