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

# 環境變數設置
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[Tuple[str, str]]) -> List[dict]:
    """將聊天歷史格式化為Grok API需要的格式"""
    formatted_messages = [
        {
            "role": "system",
            "content": """You are Grok, a chatbot inspired by the Hitchhikers Guide to the Galaxy. 
            Please respond in Traditional Chinese (繁體中文). 
            Keep your responses natural and conversational while maintaining Traditional Chinese characters."""
        }
    ]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            formatted_messages.append({
                "role": "user",
                "content": user_msg
            })
        if assistant_msg:
            formatted_messages.append({
                "role": "assistant",
                "content": assistant_msg
            })
    
    return formatted_messages

def respond(message: str, history: List[Tuple[str, str]]) -> Tuple[str, List[Tuple[str, str]]]:
    """處理用戶輸入並生成回應"""
    # 格式化歷史記錄並添加新消息
    history_for_api = format_history(history + [(message, None)])
    
    # 獲取Grok的回應
    response = call_grok_api(history_for_api)
    
    # 返回更新後的歷史記錄和空字符串(用於清除輸入框)
    history.append((message, response))
    return "", history

# 創建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("""
    # Grok 聊天機器人
    與 Elon Musk 的 Grok AI 聊天！這個聊天機器人使用了 xAI 的 Grok API。
    """)
    
    chatbot = gr.Chatbot(
        value=[],
        height=500,
        bubble_full_width=False,
        show_copy_button=True,
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="輸入訊息...",
            container=False,
            scale=7,
            autofocus=True,  # 自動聚焦到輸入框
        )
        clear = gr.ClearButton([msg, chatbot], variant="secondary", scale=1)
    
    # 注意這裡輸出添加了msg，用於清除輸入框
    msg.submit(respond, [msg, chatbot], [msg, chatbot], api_name="chat")
    
demo.queue()

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