import os
import pickle
import gradio as gr
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Load pre-processed vector data from .pkl and .faiss files
VECTOR_STORE_DIR = "./faiss"
PICKLE_PATH = os.path.join(VECTOR_STORE_DIR, "index.pkl")

# Load documents
if os.path.exists(PICKLE_PATH):
    with open(PICKLE_PATH, 'rb') as file:
        documents = pickle.load(file)
else:
    raise FileNotFoundError(f"The file {PICKLE_PATH} was not found.")

# Initialize embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Load or create FAISS vector store
FAISS_INDEX_PATH = os.path.join(VECTOR_STORE_DIR, "index.faiss")
if os.path.exists(FAISS_INDEX_PATH):
    vector_store = FAISS.load_local(VECTOR_STORE_DIR, embeddings, allow_dangerous_deserialization=True)
else:
    vector_store = FAISS.from_texts(documents, embeddings)
    vector_store.save_local(VECTOR_STORE_DIR)

# Create retriever from vector store
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})

# Define the prompt template
# prompt_template = """
# You are an AI assistant that helps answer questions based on the following context:

# {context}

# Question: {question}
# Answer:
# """

prompt_template = """
You are an AI assistant that helps answer questions based on the following context:

{context}

Question: {question}
Answer:
"""


prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=prompt_template
)

# Initialize LLM model
REPO_ID = "mistralai/Mistral-7B-Instruct-v0.2"
SEC_KEY = os.getenv("SEC_KEY")# Replace with your HuggingFace API key

model = HuggingFaceHub(
    repo_id=REPO_ID,
    huggingfacehub_api_token=SEC_KEY,
    model_kwargs={"max_length":128, "temperature":0.7}
)

# Create the RAG chain
rag_chain = LLMChain(llm=model, prompt=prompt)

def answer_question(question):
    docs = retriever.get_relevant_documents(question)
    context = "\n".join([doc.page_content for doc in docs])
    raw_output = rag_chain.run({"context": context, "question": question})

    # Attempt to isolate only the answer
    if "Answer:" in raw_output:
        final_answer = raw_output.split("Answer:")[-1].strip()
    else:
        # If 'Answer:' is not found for some reason, just strip leading/trailing whitespace
        final_answer = raw_output.strip()

    return final_answer

# Define the Gradio function
def ask_model(history, question):
    # Simulated response (Replace this with your actual model logic)
    # response = f"AI Response to: '{question}'"
    response = answer_question(question)
    history.append((question, response))  # Append user input and bot response as a tuple
    return history

# Footer content
footer_md = """
---
© 2024 by [Shashi Kiran, Karthik K, Venkara V V, Navin Kumar N, Jyoti Bavne]. All rights reserved.
This app was developed by **6505 Project Team** as part of Final Project.
For inquiries, please contact: [schandrappa@student.fairfield.edu](mailto:schandrappa@student.fairfield.edu)
"""

# Build Gradio UI
app = gr.Blocks()

with app:
    # Title and description
    gr.Markdown(
        """
        # Banking Regulations Compliance ChatBOT
        Ask questions and get responses generated by a state-of-the-art AI model on Banking related content (Outputs will be biased towards RBI (Indian Banking Regulatory) as it only trained on it)!
        """
    )

    # Main chatbot interface
    chatbot = gr.Chatbot(
        label="Chat with the Bot"
    )

    question_box = gr.Textbox(
        label="Your Message",
        placeholder="Type your message here...",
        lines=1,
    )
    submit_button = gr.Button("Send")

    # Add examples
    gr.Examples(
        examples=[
            "What is Compliance?",
            "Can you summarize the RBI Guidelines?",
            "Can you summarize the RBI Guidelines related to gold loans?",
        ],
        inputs=question_box,
    )

    # Link question input to the chatbot
    submit_button.click(ask_model, inputs=[chatbot, question_box], outputs=chatbot)

    gr.Markdown(footer_md)

app.launch(share=True)