import gradio as gr
import pickle
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline

# Step 1: Load saved chunks
with open("chunks.pkl", "rb") as f:
    chunks = pickle.load(f)

# Step 2: Load FAISS index
index = faiss.read_index("gitlab_index.faiss")


#Loading the embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

# Loading a Question-Answer model from Transformers pipeline

qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")

# Define retrieval function
def query_knowledge_base(query, top_k=3):
    query_embedding = embedding_model.encode([query])
    distances, indices = index.search(np.array(query_embedding), top_k)
    results = [chunks[i] for i in indices[0]]
    return results

def generate_answer(context, question):
    input_text = context if isinstance(context, str) else " ".join(context)
    result = qa_pipeline(question=question, context=input_text)
    return result["answer"]

def rag_chatbot(question):
    try:
        # Step 1: Embed the query and search the index
        top_chunks = query_knowledge_base(question)

        # Step 2: Combine top chunks into a single context string
        context = " ".join(top_chunks)

        result = generate_answer(context, question)

        return result

    except Exception as e:
        return f"An error occurred: {str(e)}"


# Gradio Interface

def chat_interface_fn(message, history):
    response = rag_chatbot(message)
    return response  # <-- Just return the response string


gr.ChatInterface(
    fn=chat_interface_fn,
    chatbot=gr.Chatbot(type='messages'),
    title="GitLab Handbook Assistant",
    description="Ask me about GitLab hiring or company policies!"
).launch(share=True, debug=True)