import gradio as gr
import os
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 HuggingFacePipeline
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain_community.llms import HuggingFaceEndpoint
from pathlib import Path
import chromadb
from unidecode import unidecode
from transformers import AutoTokenizer
import transformers
import torch
import tqdm 
import accelerate
import re

list_llm = ["HuggingFaceH4/zephyr-7b-beta", "mistralai/Mistral-7B-Instruct-v0.2"]
list_llm_simple = [os.path.basename(llm) for llm in list_llm]

# Load PDF document and create doc splits
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

# Create vector database
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

# Initialize langchain LLM chain
def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
    progress(0.1, desc="Initializing HF Hub...")
    if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            temperature=temperature,
            max_new_tokens=max_tokens,
            top_k=top_k,
            load_in_8bit=True,
        )
    # Add other LLM models initialization here
    progress(0.75, desc="Defining buffer memory...")
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key='answer',
        return_messages=True
    )
    retriever=vector_db.as_retriever()
    progress(0.8, desc="Defining retrieval chain...")
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm,
        retriever=retriever,
        chain_type="stuff", 
        memory=memory,
        return_source_documents=True,
        verbose=False,
    )
    progress(0.9, desc="Done!")
    return qa_chain

def format_chat_history(message, chat_history):
    formatted_chat_history = []
    for user_message, bot_message in chat_history:
        formatted_chat_history.append(f"User: {user_message}")
        formatted_chat_history.append(f"Assistant: {bot_message}")
    return formatted_chat_history

def conversation(qa_chain, message, history):
    formatted_chat_history = format_chat_history(message, history)
    response = qa_chain({"question": message, "chat_history": formatted_chat_history})
    response_answer = response["answer"]
    response_sources = response["source_documents"]
    response_source1 = response_sources[0].page_content.strip()
    response_source2 = response_sources[1].page_content.strip()
    response_source3 = response_sources[2].page_content.strip()
    response_source1_page = response_sources[0].metadata["page"] + 1
    response_source2_page = response_sources[1].metadata["page"] + 1
    response_source3_page = response_sources[2].metadata["page"] + 1
    new_history = history + [(message, response_answer)]
    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page

def upload_file(file_obj):
    list_file_path = []
    for idx, file in enumerate(file_obj):
        file_path = file_obj.name
        list_file_path.append(file_path)
    return list_file_path

def demo():
    with gr.Blocks(theme="base") as demo:
        vector_db = gr.State()
        qa_chain = gr.State()
        collection_name = gr.State()
        
        gr.Markdown(
        """<center><h2>PDF-based chatbot</center></h2>
        <h3>Ask any questions about your PDF documents</h3>""")
        
        with gr.Tab("Step 1 - Upload PDF"):
            with gr.Row():
                document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
        
        with gr.Tab("Step 2 - Process document"):
            with gr.Row():
                db_btn = gr.Radio(["ChromaDB"], label="Vector database type", value = "ChromaDB", type="index", info="Choose your vector database")
            with gr.Accordion("Advanced options - Document text splitter", open=False):
                with gr.Row():
                    slider_chunk_size = gr.Slider(minimum = 100, maximum = 1000, value=600, step=20, label="Chunk size", info="Chunk size", interactive=True)
                with gr.Row():
                    slider_chunk_overlap = gr.Slider(minimum = 10, maximum = 200, value=40, step=10, label="Chunk overlap", info="Chunk overlap", interactive=True)
            with gr.Row():
                db_progress = gr.Textbox(label="Vector database initialization", value="None")
            with gr.Row():
                db_btn = gr.Button("Generate vector database")
            
        with gr.Tab("Step 3 - Initialize QA chain"):
            with gr.Row():
                llm_btn = gr.Radio(list_llm_simple, \
                    label="LLM models", value = list_llm_simple[0], type="index", info="Choose your LLM model")
            with gr.Accordion("Advanced options - LLM model", open=False):
                with gr.Row():
                    slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.7, step=0.1, label="Temperature", info="Model temperature", interactive=True)
                with gr.Row():
                    slider_maxtokens = gr.Slider(minimum = 224, maximum = 4096, value=1024, step=32, label="Max Tokens", info="Model max tokens", interactive=True)
                with gr.Row():
                    slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k samples", info="Model top-k samples", interactive=True)
            with gr.Row():
                llm_progress = gr.Textbox(label="LLM initialization", value="None")
            with gr.Row():
                llm_btn = gr.Button("Initialize LLM chain")
            
        with gr.Tab("Step 4 - Chat"):
            with gr.Row():
                message = gr.Textbox(label="Your question")
                ask_btn = gr.Button("Ask")
            with gr.Row():
                answer = gr.Textbox(label="Answer", value="Ask your question to get an answer")
                chat_history = gr.Textbox(label="Chat history", value="Chat history")
                source1 = gr.Textbox(label="Source 1", value="Source 1")
                source2 = gr.Textbox(label="Source 2", value="Source 2")
                source3 = gr.Textbox(label="Source 3", value="Source 3")

    @demo.func
    def upload_file(file_obj):
        list_file_path = []
        for idx, file in enumerate(file_obj):
            file_path = file_obj.name
            list_file_path.append(file_path)
        return list_file_path
    
    @demo.func
    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
    
    @demo.func
    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

    @demo.func
    def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
        progress(0.1, desc="Initializing HF Hub...")
        if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
            llm = HuggingFaceEndpoint(
                repo_id=llm_model, 
                temperature=temperature,
                max_new_tokens=max_tokens,
                top_k=top_k,
                load_in_8bit=True,
            )
        # Add other LLM models initialization here
        progress(0.75, desc="Defining buffer memory...")
        memory = ConversationBufferMemory(
            memory_key="chat_history",
            output_key='answer',
            return_messages=True
        )
        retriever=vector_db.as_retriever()
        progress(0.8, desc="Defining retrieval chain...")
        qa_chain = ConversationalRetrievalChain.from_llm(
            llm,
            retriever=retriever,
            chain_type="stuff", 
            memory=memory,
            return_source_documents=True,
            verbose=False,
        )
        progress(0.9, desc="Done!")
        return qa_chain

    @demo.func
    def format_chat_history(message, chat_history):
        formatted_chat_history = []
        for user_message, bot_message in chat_history:
            formatted_chat_history.append(f"User: {user_message}")
            formatted_chat_history.append(f"Assistant: {bot_message}")
        return formatted_chat_history

    @demo.func
    def conversation(qa_chain, message, history):
        formatted_chat_history = format_chat_history(message, history)
        response = qa_chain({"question": message, "chat_history": formatted_chat_history})
        response_answer = response["answer"]
        response_sources = response["source_documents"]
        response_source1 = response_sources[0].page_content.strip()
        response_source2 = response_sources[1].page_content.strip()
        response_source3 = response_sources[2].page_content.strip()
        response_source1_page = response_sources[0].metadata["page"] + 1
        response_source2_page = response_sources[1].metadata["page"] + 1
        response_source3_page = response_sources[2].metadata["page"] + 1
        new_history = history + [(message, response_answer)]
        return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page

    # Define file upload actions
    demo.upload_file(upload_file)

    # Define document processing actions
    demo.load_doc(load_doc)
    demo.create_db(create_db)

    # Define LLM chain initialization actions
    demo.initialize_llmchain(initialize_llmchain)

    # Define conversation actions
    demo.format_chat_history(format_chat_history)
    demo.conversation(conversation)

    demo.launch()