# app.py
import os
import re
import gradio as gr
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from sentence_transformers import SentenceTransformer, util

# ----------- Load FAISS Index -----------
openai_api_key = os.getenv("OPENAI_API_KEY")

try:
    embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
    vectorstore = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
    print("✅ FAISS index loaded successfully.")
except Exception as e:
    retriever = None
    print("❌ No FAISS index found. Please upload faiss_index/ folder.", e)

# ----------- Reranker -----------
rerank_model = SentenceTransformer("all-MiniLM-L6-v2")

def rerank(question, docs, top_k=4):
    if not docs:
        return []
    query_emb = rerank_model.encode(question, convert_to_tensor=True)
    doc_embs = rerank_model.encode([doc.page_content for doc in docs], convert_to_tensor=True)
    scores = util.pytorch_cos_sim(query_emb, doc_embs)[0]
    top_results = scores.topk(top_k)
    return [docs[i] for i in top_results.indices]

# ----------- LLM Setup -----------
llm = ChatOpenAI(
    model_name="gpt-3.5-turbo",
    temperature=0,
    openai_api_key=openai_api_key
)

# ----------- Ask Function -----------
def clean_answer(answer: str) -> str:
    # Remove boilerplate or unwanted wording
    patterns = [
        r"according to the (provided )?context.*",
        r"This is mentioned in the first sentence.*",
        r"The FX Manual does not provide further details.*",
        r"This information.*FX Manual.*"
    ]
    for p in patterns:
        answer = re.sub(p, "", answer, flags=re.IGNORECASE).strip()
    return answer

def ask_question(question: str) -> str:
    if retriever is None:
        return "❌ No FAISS index found. Please upload faiss_index/ folder."
    candidate_docs = retriever.invoke(question)
    relevant_docs = rerank(question, candidate_docs, top_k=4)
    context = "\n\n".join([doc.page_content for doc in relevant_docs])
    messages = [
        {"role": "system", "content": (
            "You are an assistant specialized in the BSP FX Manual. "
            "Answer concisely and strictly based on the FX Manual excerpts provided. "
            "Cite relevant sections if available. "
            "Make the response brief and straightforward as much as possible. "
            "If the answer is not in the excerpts, say: "
            "'I could not find that in the FX Manual.'"
        )},
        {"role": "user", "content": f"FX Manual Excerpts:\n{context}\n\nQuestion: {question}"}
    ]
    response = llm.invoke(messages)
    return clean_answer(response.content)

# ----------- Gradio UI -----------
title = """
📘 FX Manual Chatbot <span style="font-size:16px; font-weight:normal;">version 1.3</span>
"""

description = """
Ask questions related to the FX Manual.  
Powered by the **International Operations Department - Policy Studies and Information Systems Group (IOD-PSISG)**  

---

### Notes:
1. The chatbot works best for simple and straightforward inquiries. If you are unsatisfied with the answer, try rewording the question.  
2. This chatbot is an AI-powered tool and may provide incomplete or inaccurate information. It should be used only as a supplementary resource. For official responses to policy inquiries, especially those that are complex or highly technical, please email **iod-ipds@bsp.gov.ph**.  
3. For technical issues or bugs that you would want to flag, kindly email **SueroJA@bsp.gov.ph**.  
"""

demo = gr.Interface(
    fn=ask_question,
    inputs=gr.Textbox(lines=2, placeholder="Ask a question related to the FX Manual..."),
    outputs="text",
    title=title,
    description=description,
)

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