import os
import re
import lz4
import nltk
import copy
import time
import vllm
import msgpack
import validators
import numpy as np
import pandas as pd
import streamlit as st
from typing import List
from numpy import ndarray
from outlines import models
from llama_cpp import Llama
import hydralit_components as hc
from transformers import AutoTokenizer
from qdrant_client import QdrantClient
from optimum_encoder import OptimumEncoder
from huggingface_hub import snapshot_download
from unstructured.partition.auto import partition
from fastembed import SparseEmbedding, SparseTextEmbedding
from unstructured.nlp.tokenize import download_nltk_packages
from scipy.sparse import csr_matrix, save_npz, load_npz, vstack
from langchain_experimental.text_splitter import SemanticChunker
from langchain_community.document_loaders import WikipediaLoader, WebBaseLoader
from qdrant_client.models import (
    NamedSparseVector,
    NamedVector,
    SparseVector,
    PointStruct,
    ScoredPoint,
    Prefetch,
    FusionQuery,
    Fusion,
    SearchRequest,
    Modifier,
    OptimizersConfigDiff,
    HnswConfigDiff,
    Distance,
    VectorParams,
    SparseVectorParams,
    SparseIndexParams
)

def make_points(texts: List[str], metadatas: List[dict], dense: List[List[float]], sparse: List[SparseEmbedding])-> List[PointStruct]:
    points = []
    for idx, (text, metadata, sparse_vector, dense_vector) in enumerate(zip(texts, metadatas, sparse, dense)):
        sparse_vec = SparseVector(indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist())
        point = PointStruct(
            id=idx,
            vector={
                "text-sparse": sparse_vec,
                "text-dense": dense_vector,
            },
            payload={
                "text": text,
                "metadata": metadata
            }
        )
        points.append(point)
    return points

def transform_query(query: str) -> str:
    """ For retrieval, add the prompt for query (not for documents).
    """
    return f'Represent this sentence for searching relevant passages: {query}'

def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding):
    dense_embeddings = dense_model.embed_query(transform_query(query))[0]
    sparse_embeddings = list(sparse_model.query_embed(query))[0]
    
    return client.query_points(
        collection_name=collection_name,
        prefetch=[
            Prefetch(query=sparse_embeddings.as_object(), using="text-sparse", limit=10),
            Prefetch(query=dense_embeddings, using="text-dense", limit=10)
        ],
        query=FusionQuery(fusion=Fusion.RRF),
        limit=3
    )

def main(query: str, client: QdrantClient, collection_name: str, llm: Llama, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding):
    scored_points = query_hybrid_search(query, client, collection_name, dense_model, sparse_model).points

    docs = [(scored_point.payload['text'], scored_point.payload['metadata']) for scored_point in scored_points]
    contents, metadatas = [list(t) for t in zip(*docs)]

    context = "\n".join(contents)
    
    seen_values = set()
    result_metadatas = "\n\n".join(
        f'{value}'
        for metadata in metadatas
        for key, value in metadata.items()
        if (value not in seen_values and not seen_values.add(value))
    )

    sampling_params_1 = vllm.SamplingParams(temperature=0, max_tokens=3000)
    prompt_1 = [
        {
            "role": "system", 
            "content": "You are a helpful assistant."
        },
        {
            "role": "user", 
            "content": st.session_state.qa_prompt(query, context)
        }
    ]
    inputs_1 = tokenizer.apply_chat_template(prompt_1, tokenize=False, add_generation_prompt=True)
    response = llm.generate(prompts=inputs_1, sampling_params=sampling_params_1)

    text = response[0].outputs[0].text
    print(f'TEXT: {response}')

    sampling_params_2 = vllm.SamplingParams(temperature=0.75, max_tokens=3000)
    prompt_2 = [
        {"role": "system", 
         "content": """Act like a professional summary writer. You have been providing summarization services for various types of documents, including academic papers, legal texts, and business reports, for over 20 years. 
         Your expertise includes extracting key points and important details concisely without adding unnecessary introductory phrases."""
        },
        {
            "role": "user",
            "content": f"""Write a summary of the following text delimited by triple backquotes. Ensure the summary covers the key points of the text. Do not introduce the summary with sentences like "Here is the summary:" or similar. The summary should be detailed, precise, and directly convey the essential information from the text.
              
            ```{text}```
              
            Let's think step-by-step."""
        }
    ]
    inputs_2 = tokenizer.apply_chat_template(prompt_2, tokenize=False, add_generation_prompt=True)
    output = llm.generate(prompts=inputs_2, sampling_params=sampling_params_2)
    
    answer = output[0].outputs[0].text
    answer_with_metadatas = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}"
    
    print(f'OUTPUT: {output}')

    if st.session_state.documents_only:
        return answer if 'no_answer' in text else answer_with_metadatas
    else:
        return f'Internal Knowledge :\n\n{answer}' if 'knowledge_topic' in text else f'Documents Based :\n\n{answer_with_metadatas}'

@st.cache_resource
def load_models_and_documents():
    with st.spinner('Load models...'):
        model_path = snapshot_download("astronomer/Llama-3-8B-Instruct-GPTQ-4-Bit")
        
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        
        llm = vllm.LLM(
            model_path, 
            tensor_parallel_size=1,
            max_model_len=8192,
            trust_remote_code=True,
            enforce_eager=True,
            quantization='gptq',
            gpu_memory_utilization=0.6,
            dtype='float16'
            #load_format='npcache'
        )
        model = models.VLLM(llm)

        dense_model = OptimumEncoder(
            device="cuda",
            cache_dir=os.getenv('HF_HOME')
        )
        
        sparse_model = SparseTextEmbedding(
            'Qdrant/bm42-all-minilm-l6-v2-attentions',
            cache_dir=os.getenv('HF_HOME'),
            providers=['CPUExecutionProvider']
        )

        download_nltk_packages()
        
    client = QdrantClient(':memory:')
    collection_name = 'collection_demo'
    
    client.create_collection(
        collection_name,
        {
            "text-dense": VectorParams(
                size=1024,
                distance=Distance.COSINE,
                on_disk=False
            )
        },
        {
            "text-sparse": SparseVectorParams(
                index=SparseIndexParams(
                    on_disk=False
                ),
                modifier=Modifier.IDF
            )
        },
        2,
        optimizers_config=OptimizersConfigDiff(
            indexing_threshold=0,
            default_segment_number=4
        ),
        hnsw_config=HnswConfigDiff(
            on_disk=False, 
            m=64,
            ef_construct=512
        )
    )
        
    with st.spinner('Parse and chunk documents...'):
        name = 'action_rpg'
        embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings')

        texts_path = os.path.join(embeddings_path, name + '_texts')
        metadatas_path = os.path.join(embeddings_path, name + '_metadatas')
        dense_path = os.path.join(embeddings_path, name + '_dense.npz')
        sparse_path = os.path.join(embeddings_path, name + '_sparse.npz')

        if not os.path.exists(embeddings_path):
            os.mkdir(embeddings_path)

            docs_1 = WikipediaLoader(query='Action-RPG').load()
            docs_2 = WikipediaLoader(query='Real-time strategy').load()
            docs_3 = WikipediaLoader(query='First-person shooter').load()
            docs_4 = WikipediaLoader(query='Multiplayer online battle arena').load()
            docs_5 = WikipediaLoader(query='List of video game genres').load()
            docs = docs_1 + docs_2 + docs_3 + docs_4 + docs_5

            texts, metadatas = [], []
            for doc in docs:
                texts.append(doc.page_content)
                del doc.metadata['title']
                del doc.metadata['summary']
                metadatas.append(doc.metadata)
            
            docs_texts, docs_metadatas, dense_embeddings, sparse_embeddings = chunk_documents(texts, metadatas, dense_model, sparse_model)

            with lz4.frame.open(texts_path, mode="wb") as outfile_texts:
                packed_texts = msgpack.packb(docs_texts, use_bin_type=True)
                outfile_texts.write(packed_texts)

            with lz4.frame.open(metadatas_path, mode="wb") as outfile_metadatas:
                packed_metadatas = msgpack.packb(docs_metadatas, use_bin_type=True)
                outfile_metadatas.write(packed_metadatas)
            
            np.savez_compressed(dense_path, *dense_embeddings)
            max_index = 0
            for embedding in sparse_embeddings:
                if embedding.indices.size > 0:
                    max_index = max(max_index, np.max(embedding.indices))
        
            sparse_matrices = []
            for embedding in sparse_embeddings:
                data = embedding.values
                indices = embedding.indices
                indptr = np.array([0, len(data)])
                matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1))
                sparse_matrices.append(matrix)
            
            combined_sparse_matrix = vstack(sparse_matrices)
            save_npz(sparse_path, combined_sparse_matrix)
        else:
            with lz4.frame.open(texts_path, mode="r") as data_file_texts:
                byte_data_texts = data_file_texts.read()
                
            with lz4.frame.open(metadatas_path, mode="r") as data_file_metadatas:
                byte_data_metadatas = data_file_metadatas.read()
                
            docs_texts = msgpack.unpackb(byte_data_texts, raw=False)
            docs_metadatas = msgpack.unpackb(byte_data_metadatas, raw=False)

            dense_embeddings = list(np.load(dense_path).values())

            sparse_embeddings = []
            loaded_sparse_matrix = load_npz(sparse_path)
        
            for i in range(loaded_sparse_matrix.shape[0]):
                row = loaded_sparse_matrix.getrow(i)
                values = row.data
                indices = row.indices
                embedding = SparseEmbedding(values, indices)
                sparse_embeddings.append(embedding)

    with st.spinner('Save documents...'):
        client.upsert(
            collection_name,
            make_points(
                docs_texts, 
                docs_metadatas,
                dense_embeddings, 
                sparse_embeddings
            )
        )
    client.update_collection(
        collection_name=collection_name,
        optimizer_config=OptimizersConfigDiff(indexing_threshold=20000)
    )       

    return client, collection_name, tokenizer, model, llm, dense_model, sparse_model

def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding):
    import time
    text_splitter = SemanticChunker(
        dense_model,
        breakpoint_threshold_type='standard_deviation'
    )
    docs = text_splitter.create_documents(texts, metadatas)
    documents, metadatas_docs = zip(*[(doc.page_content, doc.metadata) for doc in docs])

    documents = list(documents)
    metadatas_docs = list(metadatas_docs)

    start_dense = time.time()
    dense_embeddings = dense_model.embed_documents(documents)
    end_dense = time.time()
    final_dense = end_dense - start_dense
    print(f'DENSE TIME: {final_dense}')

    start_sparse = time.time()
    sparse_embeddings = list(sparse_model.embed(documents, 32))
    end_sparse = time.time()
    final_sparse = end_sparse - start_sparse
    print(f'SPARSE TIME: {final_sparse}')
    
    return documents, metadatas_docs, dense_embeddings, sparse_embeddings

def on_change_documents_only():
    st.session_state.qa_prompt = lambda query, context: (      
        f"""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, reply with 'no_answer'. Use three sentences maximum and keep the answer concise.
        
        Question: {query} 

        Context: {context} 

        Answer:"""
        if st.session_state.documents_only else 
        f"""If the context is not relevant, please answer the question by using your own knowledge about the topic.
        If you decide to provide information using your own knowledge or general knowledge, write 'knowledge_topic' at the top of your answer
        
        {context}

        Question: {query}"""
    )
    
    st.session_state.tooltip = 'The AI answer your questions only considering the documents provided' if st.session_state.documents_only else """The AI answer your questions considering the documents provided, and if it doesn't found the answer in them, try to find in its own internal knowledge"""
    
    
if __name__ == '__main__':
    st.set_page_config(page_title="Multipurpose AI Agent",layout="wide", initial_sidebar_state='collapsed')
    
    menu_id = hc.nav_bar(
        menu_definition=[
            {'id': 'ChatBot', 'label': 'ChatBot', 'ttip': 'ChatBot'},
            {'id': 'Documents', 'label': 'Documents', 'ttip': 'Documents'}
        ],
        hide_streamlit_markers=False,
        sticky_nav=True,
        sticky_mode='pinned'
    )
    st.title('Multipurpose AI Agent')
    #st.markdown("<h1 style='position: fixed; top: 0; left: 0; width: 100%; padding: 10px; text-align: left; color: black;'>Multipurpose AI Agent</h1>", unsafe_allow_html=True)

    if "qa_prompt" not in st.session_state:
        st.session_state.qa_prompt = lambda query, context: (
            f"""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, reply with 'no_answer'. Use three sentences maximum and keep the answer concise.
            
            Question: {query} 
              
            Context: {context} 
              
            Answer:"""
        )

    if "tooltip" not in st.session_state:
        st.session_state.tooltip = 'The AI answer your questions only considering the documents provided'
    
    client, collection_name, tokenizer, model, llm, dense_model, sparse_model = load_models_and_documents()
    
    if menu_id == 'Documents':
        editable_df = st.data_editor(
            pd.DataFrame([]), 
            num_rows="dynamic",
            use_container_width=True,
            column_config={
                "document": st.column_config.TextColumn(
                    "Documents",
                    help="Name of the document",
                    required=True
                ),
                "type": st.column_config.SelectboxColumn(
                    'File type',
                    help='The file format extension of this document',
                    required=True,
                    options=[
                        'DOC',
                        'DOCX',
                        'PPT',
                        'PPTX',
                        'XSLX'
                    ]
                ),
                "link": st.column_config.LinkColumn(
                    'Links',
                    help='Link to the document',
                    required=False
                ),
                "time": st.column_config.DatetimeColumn(
                    'Date and hour',
                    help='When this document has been ingested here for the last time',
                    required=True
                ),
                "toggle": st.column_config.CheckboxColumn(
                    'Enable/Disable',
                    help='Either to enable or disable the ability for the ai to find this document',
                    required=True,
                    default=True
                ) 
            }
        )

    conversations_path = os.path.join(os.getenv('HF_HOME'), 'conversations')

    try:
        with lz4.frame.open(conversations_path, mode='r') as fp:
            conversations = msgpack.unpackb(fp.read(), raw=False)
    except:
        conversations = {}

    if menu_id == 'ChatBot':
        with st.sidebar:
            if 'chat_id' not in st.session_state:
                st.session_state.chat_id = st.selectbox(
                    label='Choose a conversation',
                    options=[None] + list(conversations.keys()),
                    format_func=lambda x: conversations.get(x, 'New Chat'),
                    placeholder='_',
                )
            else:
                st.session_state.chat_id = st.selectbox(
                    label='Choose a conversation',
                    options=[None, st.session_state.chat_id] + list(conversations.keys()),
                    index=1,
                    format_func=lambda x: conversations.get(x, 'New Chat' if x != st.session_state.chat_id else st.session_state.chat_title),
                    placeholder='_',
                )
            st.session_state.chat_title = f'ChatSession-{st.session_state.chat_id}'

        chat_id_path = os.path.join(os.getenv('HF_HOME'), f'{st.session_state.chat_id}-st_messages')
        try:
            with lz4.frame.open(chat_id_path, mode='r') as fp:
                st.session_state.messages = msgpack.unpackb(fp.read(), raw=False)
        except:
            st.session_state.messages = []
            
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
                
        if prompt := st.chat_input("Message Video Game Assistant"):
            if st.session_state.chat_id not in conversations.keys():
                sampling_params = vllm.SamplingParams(temperature=0.75, max_tokens=35)
                prompt_conversation = [{"role": "user", "content": f"{prompt}\nExplain the above in one sentence:"}]
                inputs = tokenizer.apply_chat_template(prompt_conversation, tokenize=False, add_generation_prompt=True)
                outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)

                st.session_state.chat_id = outputs[0].outputs[0].text
                st.session_state.chat_title = f'ChatSession-{st.session_state.chat_id}'
                
                conversations[st.session_state.chat_id] = st.session_state.chat_title
                with lz4.frame.open(conversations_path, mode='wb') as fp:
                    packed_bytes = msgpack.packb(conversations, use_bin_type=True)
                    fp.write(packed_bytes)
            
            st.chat_message("user").markdown(prompt)
            st.session_state.messages.append({"role": "user", "content": prompt})
            print(f'PROMPT: {prompt}')
            ai_response = main(prompt, client, collection_name, llm, dense_model, sparse_model)
            with st.chat_message("assistant"):
                message_placeholder = st.empty()
                full_response = ""
                for chunk in re.split(r'(\s+)', ai_response):
                    full_response += chunk + " "
                    time.sleep(0.01)
                    message_placeholder.write(full_response + '▌')
                message_placeholder.write(re.sub('▌', '', full_response))
                                          
            st.session_state.messages.append({"role": "assistant", "content": full_response})
            with lz4.frame.open(chat_id_path, mode='wb') as fp:
                    packed_bytes = msgpack.packb(st.session_state.messages, use_bin_type=True)
                    fp.write(packed_bytes)

        with st.sidebar:
            st.divider()
            tooltip = 'The AI answer your questions only considering the documents provided'
            
            st.toggle(
                label="""Enable 'Documents-Only' Mode""", 
                value=True, 
                on_change=on_change_documents_only, 
                key="documents_only",
                help=st.session_state.tooltip
            )

            st.divider()
            
            url = st.text_input("Scrape an URL link :")
            if validators.url(url):
                docs = WebBaseLoader(url).load()
                print(f'WebBaseLoader: {docs[0].metadata}')
                
                texts, metadatas = [], []
                for doc in docs:
                    texts.append(doc.page_content)
                    del doc.metadata['title']
                    del doc.metadata['description']
                    del doc.metadata['language']
                    metadatas.append(doc.metadata)

                texts, metadatas, dense_embeddings, sparse_embeddings = chunk_documents(texts, metadatas, dense_model, sparse_model)
        
                client.upsert(
                    collection_name,
                    make_points(
                        texts,
                        metadatas,
                        dense_embeddings, 
                        sparse_embeddings
                    )
                )
                
                container = st.empty()
                container.success("URL content uploaded and ready!")
                time.sleep(2)
                container.empty()
            
            st.divider()
        
            uploaded_files = st.file_uploader("Upload a file :", accept_multiple_files=True, type=['docx', 'doc', 'pptx', 'ppt', 'xlsx', 'txt'])
            print(f'uploaded-files : {uploaded_files}')
            for uploaded_file in uploaded_files:
                print('count')
                elements = partition(file=uploaded_file,
                                     strategy='fast',
                                     skip_infer_table_types=['png', 'pdf', 'jpg', 'xls', 'xlsx', 'heic'],
                                     include_page_breaks=True
                                    )
                metadata_dict = {"source": uploaded_file.name}
                
                texts, metadatas = [], []
                for elem in elements:
                    texts.append(elem.text)
                    metadatas.append(metadata_dict)
                
                texts, metadatas, dense_embeddings, sparse_embeddings = chunk_documents(texts, metadatas, dense_model, sparse_model)

                client.upsert(
                    collection_name,
                    make_points(
                        texts,
                        metadatas,
                        dense_embeddings, 
                        sparse_embeddings
                    )
                )

                container = st.empty()
                container.success("Document content uploaded and ready!")
                time.sleep(2)
                container.empty()