import os
import zipfile
import gradio as gr
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from huggingface_hub import InferenceClient
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableParallel

# 1. Unzip stored FAISS index
if not os.path.exists("faiss_store"):
    with zipfile.ZipFile("faiss_store.zip", "r") as zip_ref:
        zip_ref.extractall(".")

# 2. Reload FAISS vector store
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vector_store = FAISS.load_local("faiss_store", embeddings, allow_dangerous_deserialization=True)

# 3. Setup retriever + RAG chain
retriever = vector_store.as_retriever(search_kwargs={"k": 3})
hf_token = os.environ.get("HF_TOKEN")
client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.2", token=hf_token)

prompt_template = """Use the following context about Narayanaswamy's (Bala) work experience to answer the question. If you don't know, say so.
Context: {context}
Question: {question}
Answer:"""

def generate_with_client(query, context):
    messages = [
        {"role": "system", "content": "You are a helpful assistant summarizing work experience."},
        {"role": "user", "content": prompt_template.format(context=context, question=query)}
    ]
    response = client.chat_completion(messages, max_tokens=200)
    return response.choices[0].message.content

rag_chain = (
    RunnableParallel({
        "context": retriever | RunnableLambda(lambda docs: "\n\n".join([doc.page_content for doc in docs])),
        "question": RunnablePassthrough()
    })
    | RunnableLambda(lambda x: generate_with_client(x["question"], x["context"]))
)

# 4. Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Know more about my work experience/portfolio...")
    chatbot = gr.Chatbot(height=400, type="messages")
    msg = gr.Textbox(placeholder="Ask about Narayanaswamy's (Bala's) work experience...")
    clear = gr.Button("Clear")

    def user(user_message, history):
        history.append({"role": "user", "content": user_message})
        return "", history

    def respond(history):
        user_message_content = history[-1]["content"]
        bot_response_content = rag_chain.invoke(user_message_content)
        history.append({"role": "assistant", "content": bot_response_content})
        return history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        respond, [chatbot], chatbot
    )
    clear.click(lambda: [], None, chatbot, queue=False)

demo.launch()