import os
import gradio as gr
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from openai import OpenAI

class CustomE5Embedding(HuggingFaceEmbeddings):
    def embed_documents(self, texts):
        texts = [f"passage: {t}" for t in texts]
        return super().embed_documents(texts)

    def embed_query(self, text):
        return super().embed_query(f"query: {text}")

embedding_model = CustomE5Embedding(model_name="intfloat/multilingual-e5-small")
db = FAISS.load_local("faiss_db", embedding_model, allow_dangerous_deserialization=True)
retriever = db.as_retriever(search_kwargs={"k": 20})

client = OpenAI()
model = "gpt-4o"

system_prompt = "你是清華大學校園資訊助理 AI，專門解答與課程、社團、場地與設施時段相關的問題。請根據資料內容，以台灣人熟悉的繁體中文提供清楚、簡潔且實用的回應。如無法在資料中找到答案，請誠實說明。"

prompt_template = """以下是關於國立清華大學課程、社團與校內設施的資料片段：
{retrieved_chunks}

使用者的提問是：{question}

請根據提供的內容回答問題，若提到課程請具體指出課程名稱、教師或時間地點；若提到社團或場地，也請回應具體活動內容或使用時段。
如查無相關資訊，請明確告知查無資料。
"""

def chat_with_rag(user_input):
    docs = retriever.get_relevant_documents(user_input)
    retrieved_chunks = "\n\n".join([doc.page_content for doc in docs])
    final_prompt = prompt_template.format(retrieved_chunks=retrieved_chunks, question=user_input)

    response = client.chat.completions.create(
        model=model,
        max_tokens=1000,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "assistant", "content": f"以下是相關資料片段：\n\n{retrieved_chunks}"},
            {"role": "user", "content": f"問題：{user_input}\n\n請根據資料，詳細列出所有相關內容，包括課程名稱、授課教師、時間地點等，必要時請使用條列式說明。若查無資料，也請清楚說明。"}
        ]
    )
    return response.choices[0].message.content

with gr.Blocks() as demo:
    gr.Markdown("# 🎓 校園資訊查詢小幫手")
    gr.Markdown("請輸入你想詢問的清華大學的問題，包含資工系課程、校園設施、社團，我會根據資料幫你解答 💬")

    chatbot = gr.Chatbot(label="📚 問答紀錄", height=600)
    msg = gr.Textbox(placeholder="例如：大一有哪些必修課？", label="❓ 問題輸入")

    def respond(message, chat_history_local):
        response = chat_with_rag(message)
        chat_history_local.append((message, response))
        return "", chat_history_local

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

demo.launch()