import gradio as gr
import torch
import os

from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.llms import HuggingFaceHub
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain.prompts import PromptTemplate
#from langchain.chains import (
#    StuffDocumentsChain, LLMChain, ConversationalRetrievalChain
#)
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Static model name
llm_name = "meta-llama/Llama-2-7b-chat-hf"

# Static file path for multiple files
static_file_paths = [
    "IRM ISO_IEC_27001_2022(en).pdf",
    #"SCF - Cybersecurity & Data Privacy Capability Maturity Model (CP-CMM) (2023.4).pdf",
    #"AG_Level1_V2.0_Final_20211210.pdf",
    #"CIS_Controls_v8_v21.10.pdf",
    #"CSF PDF v11.1.0-1.pdf",
    #"ISO_31000_2018(en)-1.pdf",
    #"OWASP Application Security Verification Standard 4.0.3-en-1.pdf",
    #"NIST.CSWP.29.ipd The NIST Cybersecurity Framework 2.0 202308-1 (1).pdf",
    #"ISO_IEC_27002_2022(en)-1.pdf",
]

# Use cuda for faster processing
device = torch.device("cuda" if torch.cuda.is_available() else "CPU")

# Load documents
#loaders = [PyPDFLoader(x) for x in static_file_paths]
#pages = []
#for loader in loaders:
#    pages.extend(loader.load())
#text_splitter = RecursiveCharacterTextSplitter(
#    chunk_size=600,
#    chunk_overlap=40,
#)
#doc_splits = text_splitter.split_documents(pages)
#embedding = HuggingFaceEmbeddings()
#vectordb = Chroma.from_documents(
#    documents=doc_splits,
#    embedding=embedding,
#)

# Load model
tokenizer = AutoTokenizer.from_pretrained(llm_name, token=os.environ['HUGGINGFACEHUB_API_TOKEN'],)
model = AutoModelForCausalLM.from_pretrained(llm_name, token=os.environ['HUGGINGFACEHUB_API_TOKEN'], torch_dtype=torch.float16)
model = model.to(device)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512, device=device, token=os.environ['HUGGINGFACEHUB_API_TOKEN'])
hf = HuggingFacePipeline(pipeline=pipe)

# Set up template and memory
template = """You are a helpful and appreciative cybersecurity expert who gives comprehensive answers using lists, step-by-step instructions and other aids. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
{context}
Question: {question}
Helpful Answer:
"""
prompt = PromptTemplate.from_template(template)
memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key='answer',
        return_messages=True
    )
retriever = vectordb.as_retriever()
qachain = ConversationalRetrievalChain.from_llm(
    hf,
    retriever=retriever,
    chain_type="stuff", 
    memory=memory,
    return_source_documents=True,
    combine_docs_chain_kwargs={
        "prompt": prompt,
    }
)

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

# Conversation with chatbot
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_source1_page = response_sources[0].metadata["page"] + 1
    response_source2_page = response_sources[1].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

def demo():
    with gr.Blocks(theme="base") as demo:
        qa_chain = gr.State(qachain)

        gr.Markdown(
        """<center><h2>Context Chatbot</center></h2>
        <h3>Ask any questions about your PDF documents, along with follow-ups</h3>
        When generating answers, it takes past questions into account (via conversational memory), and includes document references for clarity purposes.</i>
        """)

        # Conversation with chatbot
        with gr.Tab("Step 3 - Conversation with chatbot"):
            chatbot = gr.Chatbot(height=600)
            with gr.Row():
                msg = gr.Textbox(placeholder="Type message", container=True)
            with gr.Row():
                submit_btn = gr.Button("Submit")
                clear_btn = gr.ClearButton([msg, chatbot])
            with gr.Accordion("Advanced - Document references", open=False):
                with gr.Row():
                    response_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
                    response_source1_page = gr.Number(label="Page", scale=1)
                with gr.Row():
                    response_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
                    response_source2_page = gr.Number(label="Page", scale=1)

        # Preprocessing events
        #db_btn.click(initialize_database, outputs=[vector_db, db_progress])

        # Chatbot events
        submit_btn.click(conversation, \
            inputs=[qa_chain, msg, chatbot], \
            outputs=[qa_chain, msg, chatbot, response_source1, response_source1_page, response_source2, response_source2_page], \
            queue=False)
        clear_btn.click(lambda:[None,"",0,"",0], \
            inputs=None, \
            outputs=[chatbot], \
            queue=False)

    demo.queue().launch(debug=True)

if __name__ == "__main__":
    demo()