import gradio as gr
import os
# 移除 Hugging Face Inference Client 相關導入
# from huggingface_hub import InferenceClient
import google.generativeai as genai # <-- 導入 Google Generative AI 庫

# 移除 Hugging Face Token 讀取和 Client 初始化
# hf_token = os.getenv("HF_TOKEN")
# if not hf_token:
#     print("Error: Hugging Face Token (HF_TOKEN) not found...")
# client = InferenceClient("...", token=hf_token)

# 讀取 Google AI API 金鑰 Secret
# 確保您的 Secret 名稱為 GOOGLE_API_KEY
google_api_key = os.getenv("GOOGLE_API_KEY")

# 檢查金鑰是否存在
if not google_api_key:
    # 在 Gradio 應用啟動時顯示錯誤或拋出異常
    error_message = "錯誤：Google AI API Key (GOOGLE_API_KEY) Secret 未設定。請在 Hugging Face Space 設定中添加。"
    print(error_message)
    # 為了讓 Space 啟動不失敗，可以在這裡暫時顯示錯誤，但後續 API 呼叫會失敗
    # 也可以選擇直接 sys.exit(1) 讓 Space 構建或啟動失敗，提示使用者設定 Secret
    # import sys
    # sys.exit(1)

# 配置 genai 庫
genai.configure(api_key=google_api_key)

# 選擇一個支援聊天的 Gemini 模型
# 'gemini-pro' 是常用的文字聊天模型
# 模型名稱可以在 Google AI 文檔中找到
MODEL_NAME = "gemini-pro"

# System Prompt 讀取方式保持不變 (從環境變數讀取，請確保 Secret 名稱一致)
# 假設你在 Hugging Face Space Secrets 中設定了一個名為 SYSTEM_PROMPT 的 Secret。
SYSTEM_PROMPT_CONTENT = os.getenv("SYSTEM_PROMPT", "你現在是一名專業的昆蟲學系的大學教授，所以你要用專業的角度簡單跟學生回答昆蟲的知識，請盡可能使用流暢的中文回答。")


# 修改 respond 函數，使用 Gemini API 進行回覆生成
# 函數簽名與之前修改後的版本一致
def respond(
    message,
    history: list[tuple[str, str]], # history 包含之前的對話 [(user, assistant), ...]
    max_tokens, # 對應 generation_config 的 max_output_tokens
    temperature,
    top_p,
):
    # 如果金鑰不存在，不進行 API 呼叫，直接返回錯誤訊息
    if not google_api_key:
         yield error_message
         return # 結束函數執行

    # 構建發送給 Gemini API 的 messages 列表
    # Gemini API 的 generate_content 方法，當包含 history 時，通常期望格式為
    # [{'role': 'user', 'parts': ['...']}, {'role': 'model', 'parts': ['...']}, ...]
    # System Instruction 在 generate_content 中通常通過將其內容放在第一條 'user' 消息之前實現
    # 或者將其作為單獨的 system_instruction 參數（如果可用且通過 generate_content 傳遞）
    # 這裡我們採用將 system prompt 內容放在第一條實際 user 消息前的策略（如果歷史為空）
    # 或直接將 system prompt 作為一個獨立的 'user'/'model' 對話對放在歷史最前面 (方便處理 Gradio history)
    # 雖然不太符合官方 system_instruction 參數的用法，但能適應 Gradio 的 history 結構

    messages_for_gemini = []

    # 將 System Prompt 內容加入消息列表的開頭
    # 這是為了讓模型理解指令，儘管不是 Gemini 官方推薦的 system_instruction 參數
    # 我們將其作為一個獨立的 'user' 訊息放在最前面，後面跟一個 'model' 訊息作為確認
    messages_for_gemini.append({"role": "user", "parts": [SYSTEM_PROMPT_CONTENT]})
    messages_for_gemini.append({"role": "model", "parts": ["好的，我明白了，請開始。"]}) # 假裝模型確認了指令


    # 將 Gradio 提供的 history 轉換為 Gemini API 的格式
    # Gradio history 是 [(user_msg, bot_msg), ...]
    for user_msg, bot_msg in history:
        # 確保訊息不為 None 或空字串，避免 API 錯誤
        if user_msg:
            messages_for_gemini.append({"role": "user", "parts": [user_msg]})
        if bot_msg:
            messages_for_gemini.append({"role": "model", "parts": [bot_msg]})

    # 加入當前使用者訊息
    if message:
        messages_for_gemini.append({"role": "user", "parts": [message]})


    response = ""
    try:
        # 呼叫 Gemini API
        # 使用 generate_content 方法
        model = genai.GenerativeModel(MODEL_NAME)

        # 設定生成參數
        generation_config = genai.types.GenerationConfig(
            max_output_tokens=max_tokens, # 對應 max_tokens
            temperature=temperature,
            top_p=top_p
            # stop_sequences=[] # 如果需要，可以在這裡設定停止標記
        )

        # 呼叫 generate_content，傳入格式化好的消息列表和生成參數
        # 對於 stream=True，它返回一個 Response 迭代器
        stream = model.generate_content(
            messages_for_gemini, # <-- 傳入構建好的消息列表
            generation_config=generation_config,
            stream=True,
        )

        # 處理流式回覆
        # 迭代 Response 迭代器，獲取 chunks
        for chunk in stream:
            # chunk.text 包含當前的文本片段
            # 需要確保 chunk.text 不是 None 或空字串
            if hasattr(chunk, 'text') and chunk.text:
                 response += chunk.text
                 yield response # 逐步返回結果

    except Exception as e:
        print(f"Error during Gemini inference: {e}")
        # 顯示更詳細的錯誤信息給使用者
        yield f"生成回覆時發生錯誤: {e}"


# Gradio ChatInterface 設定保持不變
# additional_inputs 移除 System message 文字框
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    # Chatbot 组件可以考虑设置 type='messages' 并在 respond 中返回对应的格式，以消除 UserWarning
    # 例如 bot=gr.Chatbot(type='messages')
    # 这需要调整 respond 函数返回值的格式，超出当前主要任务范围，可以之后优化
)

if __name__ == "__main__":
    demo.launch()