#app.py
import time
import gradio as gr
import google.generativeai as genai
import os

# 從 Hugging Face secrets 中讀取 API 金鑰（如果需要）
api_key = os.getenv('GOOGLE_API_KEY')
if not api_key:
    raise ValueError("請設置 'GOOGLE_API_KEY' 環境變數")

# 設定 API 金鑰
genai.configure(api_key=api_key)

# 初始化模型
try:
    model = genai.GenerativeModel('gemini-1.5-pro')
    chat = model.start_chat(history=[])
    print("模型載入成功。")
except Exception as e:
    raise ValueError(f"無法載入模型：{e}")

# 將 Gradio 的歷史紀錄轉換為 Gemini 格式
def transform_history(history):
    new_history = []
    for chat in history:
        new_history.append({"parts": [{"text": chat[0]}], "role": "user"})
        new_history.append({"parts": [{"text": chat[1]}], "role": "model"})
    return new_history

# 回應生成函數
def response(message, history):
    global chat
    # 將 Gradio 的歷史紀錄轉換為 Gemini 的格式
    chat.history = transform_history(history)

    prompt = "你是一個半導體智慧製造專家，熟悉半導體製程及生產過程，請用專業的角度回應接下來的問題!"
    if len(history) == 0:
        prompt_message = f"{prompt} \n {message}"
    else:
        prompt_message = message
        
    # 發送訊息到 Gemini API
    response = chat.send_message(prompt_message)
    response.resolve()

    # 逐字回傳生成的文字，實現打字機效果
    for i in range(len(response.text)):
        time.sleep(0.05)  # 每個字符間隔 0.05 秒
        yield response.text[: i+1]
        

# 建立 Gradio 聊天界面
gr.ChatInterface(response,
                 title='智慧製造小幫手',
                 textbox=gr.Textbox(placeholder="Question to Gemini")).launch(share=True)