import openai
import gradio as gr
import os

client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])

# 檔案搜尋工具結構(schema)
file_search_tool = {
    "type": "file_search", # 工具類型
    "vector_store_ids": [
        "vs_692148c1f7148191ac1a319f4b918206" # 第一步中建立的資料庫ID
    ], # 可查詢的資料庫id
    "max_num_results": 2, # 搜尋最大資料筆數
}

def normalize_messages(messages):
    """ Gradio 訊息格式改成OpenAI """
    normalized = []

    for msg in messages:
        role = msg.get("role")
        contents = msg.get("content", [])

        new_contents = []

        # 決定目標 type
        if role in ["user", "system"]:
            target_type = "input_text"
        elif role == "assistant":
            target_type = "output_text"
        else:
            target_type = "input_text"  # fallback

        # 處理 content 陣列
        for c in contents:
            text = c.get("text", "")
            new_contents.append({
                "type": target_type,
                "text": text
            })

        normalized.append({
            "role": role,
            "content": new_contents
        })

    return normalized

def get_system_prompt():
    # 設定角色定位，與回覆方式
    system_prompt = f"""
    你是一個專業的助手可以回答有關台灣大學的行事曆與校規。
    但你可以根據以下的參考資料來回答，不使用你的內部知識，也不捏造資訊。
    可以適當地根據參考資料推論使用者問題
    一步一步思考 最後再回答
    給使用者的內容加上查詢到的條文或法規
    """
    return system_prompt

def get_response(message, history):
    """
    利用OpenAI回覆訊息
    :param message: 輸入訊息
    :param history: 歷史紀錄
    """
    system_prompt = get_system_prompt()

    # 將歷史紀錄 + system prompt + 新的輸入 當作prompt
    prompt = history + [
        {"role": "system", "content": [{"type": "input_text", "text": system_prompt}]},
        {"role": "user", "content": [{"type": "input_text", "text": message}]}
    ]
    prompt = normalize_messages(prompt)

    for p in prompt:
        print(p)
    try:
        stream = client.responses.create(
            model="gpt-5.1", # gpt-5-mini, gpt-4o-mini, gpt-4.1-mini,gpt-5-nano,gpt-5,gpt-5.1
            input=prompt,
            tools=[file_search_tool], # 設定工具
            stream=True # 串流模式
        )
        buffer = ""
        for event in stream:
            if event.type == "response.output_text.delta":
                # 逐 token（或字串片段）輸出
                buffer += event.delta or ""
                yield buffer
            elif event.type == "response.error":
                # 串流中途發生錯誤事件
                yield "Unknown streaming error"
            elif event.type == "response.completed":
                # 串流結束
                break
    except Exception as e:
        yield f"\n[其他錯誤] {e}\n"

# Gradio 的聊天介面 (ChatInterface)
# fn=get_response : 指定使用者輸入後，呼叫 get_response 這個函式來產生回覆
web_chat = gr.ChatInterface(
    fn=get_response,
)

# 主程式
if __name__ == "__main__":
    # 啟動 Gradio 服務
    # debug=True : 開啟除錯模式，可以在 console 看到更詳細的錯誤訊息
    web_chat.launch(debug=True)