import gradio as gr
import pickle
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain_together import Together


# Function to load embeddings
def load_embeddings():
    try:
        embeddings = HuggingFaceEmbeddings(
            model_name="nomic-ai/nomic-embed-text-v1",
            model_kwargs={"trust_remote_code": True, "revision": "289f532e14dbbbd5a04753fa58739e9ba766f3c7"}
        )
        print("Embeddings loaded successfully.")
        return embeddings
    except Exception as e:
        raise RuntimeError(f"Error loading embeddings: {e}")


embeddings = load_embeddings()


# Load FAISS index and pickle file without explicitly defining paths
def load_db():
    try:
        db = FAISS.load_local("law_vector_db", embeddings, allow_dangerous_deserialization=True)
        print(f"FAISS index loaded successfully.")

        with open('law_vector_db/index.pkl', 'rb') as pkl_file:
            metadata = pickle.load(pkl_file)
        print("Pickle file loaded successfully.")

        return db, metadata
    except Exception as e:
        raise RuntimeError(f"Error loading FAISS index or pickle file: {e}")


db, metadata = load_db()

# Create a retriever from the FAISS index
db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 4})

# Define the prompt template for the chatbot
prompt_template = """
<s>[INST]This is a chat template and As a legal chatbot specializing in Indian Penal Code queries, your primary objective is to provide accurate and concise information based on the user's questions. 
Do not generate your own questions and answers. You will adhere strictly to the instructions provided, offering relevant context from the knowledge base while avoiding unnecessary details. 
Your responses will be brief, to the point, and in compliance with the established format. 
If a question falls outside the given context, you will refrain from utilizing the chat history and instead rely on your own knowledge base to generate an appropriate response. 
You will prioritize the user's query and refrain from posing additional questions. 
The aim is to deliver professional, precise, and contextually relevant information pertaining to the Indian Penal Code.

CONTEXT: {context}
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:</s>[INST]
"""

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

# Together API Key (hardcoded)
TOGETHER_AI_API = "66bd7a6dc11956ddb311b773c0deabda8870e8c90e9f548ce064880ac47c4b05"

# Initialize LLM (Together API)
llm = Together(
    model="mistralai/Mistral-7B-Instruct-v0.2",
    temperature=0.5,
    max_tokens=1024,
    together_api_key=TOGETHER_AI_API
)


# Function to process user input and generate responses
def ask_question(user_question, chat_history=[]):
    try:
        # Retrieve relevant documents from FAISS index
        context_docs = db_retriever.get_relevant_documents(user_question)
        context = "\n".join(
            [doc.page_content for doc in context_docs]) if context_docs else "No relevant context found."

        # Prepare input for the model
        input_data = {
            "context": context,
            "question": user_question,
            "chat_history": "\n".join(chat_history)  # Chat history as string
        }

        # Generate the answer using Together API
        response = llm(prompt.format(**input_data))
        return response
    except Exception as e:
        return f"Error: {e}"


# Function to manage conversation flow
def chat_bot_interface(user_message, chat_history=[]):
    if not user_message:
        return chat_history, chat_history  # No update if message is empty

    # Append user message
    chat_history.append(("User", user_message))

    # Get system response
    response = ask_question(user_message, [msg[1] for msg in chat_history if msg[0] == "User"])

    # Append system response
    chat_history.append(("Assistant", response))

    return chat_history, chat_history


# Set up Gradio interface with a professional chatbot UI
with gr.Blocks() as iface:
    gr.Markdown("<h1 style='text-align: center;'>Legal Chatbot</h1>")

    chatbot = gr.Chatbot(label="Chatbot Interface")
    user_input = gr.Textbox(label="Ask a Question", placeholder="Type your question here...", lines=1)
    clear_button = gr.Button("Clear")

    # Maintain chat history
    chat_history = gr.State([])


    def clear_chat():
        return [], []


    user_input.submit(chat_bot_interface, inputs=[user_input, chat_history], outputs=[chatbot, chat_history])
    clear_button.click(clear_chat, outputs=[chatbot, chat_history])

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()