#!/usr/bin/env python3
import subprocess
 
command = 'CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python --force-reinstall --upgrade --no-cache-dir'
process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
output, error = process.communicate()
output = output.decode("utf-8")
error = error.decode("utf-8")
print("Output:", output)
print("Error:", error)
 
command2 = 'pip install langchain'
process2 = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
output2, error2 = process2.communicate()
output2 = output2.decode("utf-8")
error2 = error2.decode("utf-8")
print("Output:", output2)
print("Error:", error2)


import datetime
import glob
import json
import logging
import os
import shutil
import sys
import uuid
from json import JSONDecodeError
from multiprocessing import Pool
from pathlib import Path
from time import sleep
from typing import List, Optional

import pandas as pd
# import qdrant_client
import streamlit as st
from dotenv import load_dotenv
from langchain import LLMChain, PromptTemplate
from langchain.chains import RetrievalQA, RetrievalQAWithSourcesChain
from langchain.docstore.document import Document
from langchain.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    PDFMinerLoader,
    TextLoader,
    UnstructuredEmailLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredODTLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader,
)
from langchain.llms import CTransformers
from langchain.embeddings import HuggingFaceEmbeddings, SentenceTransformerEmbeddings
from langchain.llms import GPT4All, LlamaCpp
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma, Qdrant
from markdown import markdown
# from qdrant_client.models import Distance, VectorParams
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
# from qdrant_client import QdrantClient
# from qdrant_client.http.models import Distance, VectorParams
from tqdm import tqdm
from langchain_core.prompts import PromptTemplate
# from langchain.callbacks.base import CallbackManager
# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from tqdm.auto import tqdm

from constants import CHROMA_SETTINGS

# from constants import CHROMA_SETTINGS

load_dotenv()
######################################## CONSTANTS #########################################
# index_name = "openai-ada-002-index-1536"
# reader_model = "gpt-3.5-turbo"
# embed_model = "text-embedding-ada-002"
# embedding_dim = 768
FILE_UPLOAD_PATH = "./data/uploads/"
qdrant_dir = "./data/qdrant_storage"
# NAME_SPACE = "qademo"
os.makedirs(FILE_UPLOAD_PATH, exist_ok=True)
os.makedirs(qdrant_dir, exist_ok=True)


# Load environment variables
persist_directory = os.environ.get('PERSIST_DIRECTORY', "vector_db")
source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
embeddings_model_name = os.environ.get('EMBEDDINGS_MODEL_NAME',"all-MiniLM-L6-v2")
embeddings_dim = os.environ.get('EMBEDDINGS_DIM',384)

chunk_size = 100
chunk_overlap = 20


model_type = os.environ.get('MODEL_TYPE',"LlamaCpp")
model_path = os.environ.get('MODEL_PATH', 'models/openhermes-2.5-mistral-7b-16k.Q4_K_M.gguf')
model_n_ctx = os.environ.get('MODEL_N_CTX',32000)
reset_index = os.environ.get("RESET_INDEX",False)
collection_name = os.environ.get('COLELCTION_NAME', "my_collection")

QDRANT_HOST = os.environ.get("QDRANT_HOST", "localhost")
QDRANT_PORT = os.environ.get("QDRANT_PORT", 6333)
target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4))


print("Current working directory:", os.getcwd())
try:
    print("Contents of /app/models:", os.listdir('./models'))
except Exception as e:
    print(f"Exception occurred: {e}")


print(f"Working with model_type: {model_type} model_path: {model_path} model_n_ctx: {model_n_ctx} reset_index: {reset_index} collection_name: {collection_name}")
###########################################################################################

# Custom document loaders
class MyElmLoader(UnstructuredEmailLoader):
    """Wrapper to fallback to text/plain when default does not work"""

    def load(self) -> List[Document]:
        """Wrapper adding fallback for elm without html"""
        try:
            try:
                doc = UnstructuredEmailLoader.load(self)
            except ValueError as e:
                if 'text/html content not found in email' in str(e):
                    # Try plain text
                    self.unstructured_kwargs["content_source"]="text/plain"
                    doc = UnstructuredEmailLoader.load(self)
                else:
                    raise
        except Exception as e:
            # Add file_path to exception message
            raise type(e)(f"{self.file_path}: {e}") from e

        return doc


# Map file extensions to document loaders and their arguments
LOADER_MAPPING = {
    ".csv": (CSVLoader, {}),
    # ".docx": (Docx2txtLoader, {}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".enex": (EverNoteLoader, {}),
    ".eml": (MyElmLoader, {}),
    ".epub": (UnstructuredEPubLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".odt": (UnstructuredODTLoader, {}),
    ".pdf": (PDFMinerLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
    # Add more mappings for other file extensions and loaders as needed
}


def load_single_document(file_path: str) -> Document:
    ext = "." + file_path.rsplit(".", 1)[-1]
    if ext in LOADER_MAPPING:
        loader_class, loader_args = LOADER_MAPPING[ext]
        loader = loader_class(file_path, **loader_args)
        return loader.load()[0]

    raise ValueError(f"Unsupported file extension '{ext}'")

@st.cache_resource
def get_embedding_model():
    model_kwargs = {'device': 'cpu'}
    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name, model_kwargs=model_kwargs)
    return embeddings

print("loading the embeddings")
embeddings = get_embedding_model()

@st.cache_resource()
def get_vector_db():
    # client = qdrant_client.QdrantClient(
    # path=persist_directory, prefer_grpc=True
    # )
    # client = qdrant_client.QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
    # available_collections = client.get_collections()
    # print(f"Available collections: {available_collections}")
    # if collection_name not in available_collections:
    #     # if reset_index:
    #     #     print(f"Deleting collection and creating again")
    #     # client.delete_collection(collection_name="{collection_name}")
    #     print(f"Creating collection: {collection_name}")
    #     client.recreate_collection(
        
    #             collection_name=collection_name,
    #             vectors_config=VectorParams(size=embeddings_dim, distance=Distance.COSINE),
    #         )
    # vectordb = Qdrant(
    #     client=client, collection_name=collection_name, 
    #     embeddings=embeddings
    # )
    
    vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    return vectordb
    
print("loading the vector DB")
vectordb = get_vector_db()
print("loading the vector as retriever")
retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"k": 5})

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


@st.cache_resource 
def get_chain():
    # Callbacks support token-wise streaming
    # callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    # Verbose is required to pass to the callback manager

    # Make sure the model path is correct for your system!
    # model_name = os.getenv("MODEL_PATH", "wizardLM-7B.ggml.q4_2.bin")
    # llm = LlamaCpp(
    #     model_path=model_name, n_ctx=1024,verbose=True, n_threads=4, n_batch=512
    # )
    callbacks = []
    
    # vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
    # https://qdrant.tech/documentation/concepts/collections/
    
    print("loading the LLM model")
    match model_type:
        case "LlamaCpp":
            llm = LlamaCpp(model_path=model_path,temperature=0.1, n_gpu_layers= 32, n_gqa=8,
                            max_new_tokens=512,context_window=2048,  n_ctx=model_n_ctx, callbacks=callbacks, verbose=True)
            # llm = CTransformers(model=model_path, config={'max_new_tokens': model_n_ctx, 'temperature': 0.01,'context_length': model_n_ctx})
        case "GPT4All":
            llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
    # qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever,  return_source_documents= True)
    print("loading the QA pipeline")
    
    qa = RetrievalQAWithSourcesChain.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, reduce_k_below_max_tokens=True, return_source_documents= True, verbose=True)
    
    template = """
        
         <s> [INST]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.
        Use three sentences maximum and keep the answer as concise as possible.
        Always say "thanks for asking!" at the end of the answer.  [/INST] </s>
        
        Context: {context}
        
         [INST] Question: {question}
        
     Answer: [/INST]"""
    custom_rag_prompt = PromptTemplate.from_template(template)
    
    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | custom_rag_prompt
        | llm
        | StrOutputParser()
    )
    

    return qa,  vectordb, rag_chain


qa, vectordb, rag_chain = get_chain()

def query(question, primer, top_k_retriever):
     # Get the answer from the chain
    print("Querying the model")
    # res = qa(question)
    retrieved_docs = retriever.invoke(question)
    if len(retrieved_docs) == 0:
        return  {"answer": "No files found", "filenames" : [] }
        
    print(f"Retrieved examples: {print(retrieved_docs[0].page_content)}")
    context = "\n\n".join(doc.page_content for doc in retrieved_docs)
    prompt = f"""
       <s> [INST] You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise. [/INST] </s> 
        [INST] Question: {question} 
        Context: {context} 
        Answer: [/INST]
        
    """
    print(F"Executing Prompt: {prompt}")
    answer = rag_chain.invoke("What is Task Decomposition?")
    print(f"Result from the model: {answer}")
    # answer, docs = res.get('result', ""), res.get('source_documents', "")
    docs = [doc.metadata["source"] for doc in retrieved_docs]
    res = {"answer": answer, "filenames" : docs }
    return res


def set_state_if_absent(key, value):
    if key not in st.session_state:
        st.session_state[key] = value


# Adjust to a question that you would like users to see in the search bar when they load the UI:
DEFAULT_QUESTION_AT_STARTUP = os.getenv("DEFAULT_QUESTION_AT_STARTUP",  "What is the state of generative ai in 2022?")
DEFAULT_ANSWER_AT_STARTUP = os.getenv(
    "DEFAULT_ANSWER_AT_STARTUP",
    "",
)
DEFAULT_PRIMER = os.getenv("DEFAULT_PRIMER", f"""You are Q&A bot that answers
    user questions based on the information provided by the user above
    each question. If the information can not be found in the information
    provided by the user you truthfully say "I don't know".
    """)

# Sliders
DEFAULT_DOCS_FROM_RETRIEVER = int(os.getenv("DEFAULT_DOCS_FROM_RETRIEVER", "3"))


# st.set_page_config(
#     page_title="Open AI Demo", page_icon="https://haystack.deepset.ai/img/HaystackIcon.png"
# )

# Persistent state
set_state_if_absent("question", DEFAULT_QUESTION_AT_STARTUP)
set_state_if_absent("answer", DEFAULT_ANSWER_AT_STARTUP)
set_state_if_absent("results", None)
set_state_if_absent("primer", DEFAULT_PRIMER)

# Small callback to reset the interface in case the text of the question changes
def reset_results(*args):
    st.session_state.answer = None
    st.session_state.results = None
    st.session_state.raw_json = None


# Title
st.write("# Open LLM Semantic Search Demo")
st.markdown(
    """
This demo takes its data from PDF, DOCX, and TXT files. \n
Ask any question on this indexed data and see if OpenAI can find the correct answer to your query! \n
*Note: do not use keywords, but full-fledged questions.* The demo is not optimized to deal with keyword queries and might misunderstand you.
""",
    unsafe_allow_html=True,
)

# Sidebar
st.sidebar.header("Options")
st.sidebar.write("## File Upload:")


with st.sidebar.form("my-form", clear_on_submit=True):
    data_files = st.file_uploader(
        "Upload",
        type=["pdf", "txt", "docx","html"],
        accept_multiple_files=True,
        label_visibility="hidden",
    )
    submitted = st.form_submit_button("UPLOAD!")
    if submitted and data_files is not None:
        st.write("UPLOADED!")
        ALL_FILES = []
        META_DATA = []
        upload_dir = Path(FILE_UPLOAD_PATH) / f"{uuid.uuid4().hex}/"
        os.makedirs(upload_dir, exist_ok=True)
        for data_file in data_files:
            # Upload file
            if data_file:
                file_path = upload_dir / data_file.name
                with open(file_path, "wb") as f:
                    f.write(data_file.getbuffer())
                ALL_FILES.append(str(file_path))
                st.sidebar.write(str(data_file.name) + " &nbsp;&nbsp; ✅ ")
                META_DATA.append({"filename": data_file.name})

        if len(ALL_FILES) > 0:
        
            st.sidebar.write("Starting the document indexing ... ")
            with st.spinner(
                    "🧠 &nbsp;&nbsp; Performing indexing of uploaded documents... \n "
                ):
                
                with Pool(processes=os.cpu_count()) as pool:
                    results = []
                    with tqdm(total=len(ALL_FILES), desc='Loading new documents', ncols=80) as pbar:
                        for i, doc in enumerate(pool.imap_unordered(load_single_document, ALL_FILES)):
                            results.append(doc)
                            pbar.update()
                print(f"Loaded {len(results)} new documents from Upload")
                text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
                docs = text_splitter.split_documents(results)
                print(f"Split into {len(docs)} chunks of text (max. {chunk_size} tokens each)")
                vectordb.add_documents(docs)
                vectordb.persist()
                            
            
            st.sidebar.write("Document indexing completed &nbsp;&nbsp; ✅ ")
            ALL_FILES = []
            META_DATA = []


top_k_retriever = st.sidebar.slider(
    "Max. number of documents from retriever",
    min_value=1,
    max_value=10,
    value=DEFAULT_DOCS_FROM_RETRIEVER,
    step=1,
    on_change=reset_results,
)
# primer = st.text_input(
#     value=st.session_state.primer,
#     max_chars=1000,
#     label="primer",
# )
primer=""
question = st.text_input(
    value=st.session_state.question,
    max_chars=200,
    on_change=reset_results,
    label="question",
    label_visibility="hidden",
)
col1, col2 = st.columns(2)
col1.markdown("<style>.stButton button {width:100%;}</style>", unsafe_allow_html=True)
col2.markdown("<style>.stButton button {width:100%;}</style>", unsafe_allow_html=True)

# Run button
run_pressed = col1.button("Run")
if run_pressed:
    run_query = run_pressed or question != st.session_state.question
    # Get results for query
    if run_query and question:
        reset_results()
        st.session_state.question = question

        with st.spinner("🧠 &nbsp;&nbsp; Performing neural search on documents... \n "):
            try:
                print(f"Running the query : {question}")
                st.session_state.results = query(question, primer, top_k_retriever=top_k_retriever)
            except JSONDecodeError as je:
                st.error(
                    "👓 &nbsp;&nbsp; An error occurred reading the results. Is the document store working?"
                )
            except Exception as e:
                logging.exception(e)
                if "The server is busy processing requests" in str(e) or "503" in str(e):
                    st.error("🧑‍🌾 &nbsp;&nbsp; All our workers are busy! Try again later.")
                else:
                    st.error(f"🐞 &nbsp;&nbsp; An error occurred during the request. {str(e)}")


if st.session_state.results:
    st.write("## Results:")
    answer = st.session_state.results["answer"]
    # Hack due to this bug: https://github.com/streamlit/streamlit/issues/3190
    try:
        filenames = st.session_state.results["filenames"]
        st.write(
            markdown(f"**Answer:** \n {answer} \n\n **Using data from files**:  {filenames} \n "),
            unsafe_allow_html=True,
        )
    except Exception as e:
        st.write(
            markdown(f"Failed to find answer:"),
            unsafe_allow_html=True,
        )