import streamlit as st
from langchain_community.llms import HuggingFaceHub
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os

import zipfile


# Unzip the data folder if not already extracted
zip_path = "data.zip"
extract_folder = "data/"

if os.path.exists(zip_path) and not os.path.exists(extract_folder):
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_folder)
    print("Data folder unzipped successfully.")

import pkg_resources
installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
print(installed_packages)


# Initialize LLM
def initialize_llm():
    llm = HuggingFaceHub(
        repo_id="meta-llama/Llama-2-7b-chat-hf",
        model_kwargs={"temperature": 0.5, "max_length": 512}
    )
    return llm

# Create vector DB
def create_vector_db():
    loader = DirectoryLoader("data/", glob="*.pdf", loader_cls=PyPDFLoader)
    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = text_splitter.split_documents(documents)
    embeddings = HuggingFaceBgeEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_db = Chroma.from_documents(texts, embeddings, persist_directory="./chroma_db")
    vector_db.persist()
    return vector_db

# Setup QA Chain
def setup_qa_chain(vector_db, llm):
    retriever = vector_db.as_retriever()
    prompt_template = """You are a compassionate mental health chatbot. Respond thoughtfully:
    {context}
    User: {question}
    Chatbot:"""
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        chain_type_kwargs={"prompt": PROMPT}
    )
    return qa_chain

# Streamlit UI
st.title("🧠 Mental Health Chatbot 🤖")
st.write("A compassionate chatbot designed to assist with mental well-being.")

llm = initialize_llm()

db_path = "chroma_db"
if not os.path.exists(db_path):
    vector_db = create_vector_db()
else:
    embeddings = HuggingFaceBgeEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_db = Chroma(persist_directory=db_path, embedding_function=embeddings)

qa_chain = setup_qa_chain(vector_db, llm)

user_input = st.text_input("You: ", "")
if st.button("Send"):
    if user_input:
        response = qa_chain.run(user_input)
        st.write(f"Chatbot: {response}")
    else:
        st.warning("Please enter a valid input.")

st.markdown("**Note:** For urgent mental health concerns, contact a licensed professional.")