import streamlit as st
import os
from pathlib import Path
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFaceEndpoint
from langchain.memory import ConversationBufferMemory
from unidecode import unidecode
import chromadb
import re

list_llm = [
    "mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", 
    "mistralai/Mistral-7B-Instruct-v0.1", "google/gemma-7b-it", "google/gemma-2b-it", 
    "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", 
    "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", 
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", 
    "google/flan-t5-xxl"
]

def load_doc(list_file_path, chunk_size, chunk_overlap):
    loaders = [PyPDFLoader(x) for x in list_file_path]
    pages = []
    for loader in loaders:
        pages.extend(loader.load())
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    doc_splits = text_splitter.split_documents(pages)
    return doc_splits

def create_db(splits, collection_name):
    embedding = HuggingFaceEmbeddings()
    new_client = chromadb.EphemeralClient()
    vectordb = Chroma.from_documents(
        documents=splits,
        embedding=embedding,
        client=new_client,
        collection_name=collection_name
    )
    return vectordb

def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
    llm = HuggingFaceEndpoint(repo_id=llm_model, temperature=temperature, max_new_tokens=max_tokens, top_k=top_k)
    memory = ConversationBufferMemory(memory_key="chat_history", output_key='answer', return_messages=True)
    retriever = vector_db.as_retriever()
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm,
        retriever=retriever,
        chain_type="stuff",
        memory=memory,
        return_source_documents=True,
        verbose=False
    )
    return qa_chain

def create_collection_name(file_path):
    collection_name = Path(file_path).stem
    collection_name = unidecode(collection_name)
    collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)
    collection_name = collection_name[:50]
    if len(collection_name) < 3:
        collection_name = collection_name + 'xyz'
    if not collection_name[0].isalnum():
        collection_name = 'A' + collection_name[1:]
    if not collection_name[-1].isalnum():
        collection_name = collection_name[:-1] + 'Z'
    return collection_name

def main():
    st.title("PDF-based Chatbot")

    uploaded_files = st.file_uploader("Upload PDF documents (single or multiple)", type="pdf", accept_multiple_files=True)

    if uploaded_files:
        chunk_size = st.slider("Chunk size", min_value=100, max_value=1000, value=600, step=20)
        chunk_overlap = st.slider("Chunk overlap", min_value=10, max_value=200, value=40, step=10)

        list_file_path = [file.name for file in uploaded_files]

        if st.button("Generate Vector Database"):
            st.text("Loading documents...")
            doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
            st.text("Creating vector database...")
            collection_name = create_collection_name(list_file_path[0])
            vector_db = create_db(doc_splits, collection_name)

            llm_model = st.selectbox("Choose LLM Model", list_llm)
            temperature = st.slider("Temperature", min_value=0.01, max_value=1.0, value=0.7, step=0.1)
            max_tokens = st.slider("Max Tokens", min_value=224, max_value=4096, value=1024, step=32)
            top_k = st.slider("Top-K Samples", min_value=1, max_value=10, value=3, step=1)

            if st.button("Initialize QA Chain"):
                st.text("Initializing QA chain...")
                qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db)

                st.header("Chatbot")
                message = st.text_input("Type your message")
                if st.button("Submit"):
                    st.text("Generating response...")
                    response = qa_chain({"question": message, "chat_history": []})
                    st.write("Assistant:", response["answer"])

if __name__ == "__main__":
    main()