import gradio as gr
from openai import OpenAI
import os

# OpenRouter Client 設定
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENROUTER_API_KEY"),
)

# 自訂知識庫（key = 主題關鍵字, value = 資料內容）
knowledge_base = {
    "午餐": "我們學校提供兩輪午餐時間。第一輪是12:15-12:45，第二輪是12:45-13:15。菜單每週更新，可在校網查閱。",
    "課外活動": "我們的課外活動包括足球、籃球、辯論、機械人編程等，活動時間是放學後 3:30 至 5:00。",
    "圖書館": "圖書館平日開放時間為早上8:30至下午4:30。借書需使用學生證，每次最多借3本，期限為兩週。",
    "校車": "校車每日早上7:30從主要地鐵站出發，下午放學後15分鐘內發車，請準時到指定上車地點。",
}

# 根據輸入問題比對知識庫，回傳對應資料串成 prompt
def search_kb(question: str) -> str:
    matched_info = []
    for keyword, content in knowledge_base.items():
        if keyword in question:
            matched_info.append(f"{keyword}：{content}")
    if matched_info:
        return "以下是你可能需要的學校資料：\n" + "\n".join(matched_info)
    else:
        return "找不到相關資料，請直接由 Chatbot 回答。"

# Chatbot 回覆邏輯：將知識庫資訊加入 system prompt
def chatbot_reply(user_input, history=None):
    kb_prompt = search_kb(user_input)

    try:
        response = client.chat.completions.create(
            model="mistralai/mistral-small-3.1-24b-instruct:free",
            messages=[
                {
                    "role": "system",
                    "content": f"你是一位校園小老師，根據以下學校資料回答問題：\n{kb_prompt}"
                },
                {
                    "role": "user",
                    "content": user_input
                }
            ],
            extra_headers={
                "HTTP-Referer": "https://your-education-workshop.com",
                "X-Title": "AI Chatbot with Knowledge Base",
            },
            max_tokens=300
        )
        reply = response.choices[0].message.content
        return reply.strip()

    except Exception as e:
        return f"發生錯誤：{str(e)}"

# Gradio Web Chat App
gr.ChatInterface(
    fn=chatbot_reply,
    title="校園小老師 Chatbot (知識庫版)",
    description="請輸入關於學校生活的問題（如：課外活動、午餐、圖書館、校車），AI 小老師會回答你！",
    theme="soft"
).launch()