#list of librarys for requirement.txt
import os
import re
import hashlib

from langchain.document_loaders import PyPDFLoader

# Import embeddings module from langchain for vector representations of text
from langchain.embeddings import HuggingFaceEmbeddings

# Import text splitter for handling large texts
from langchain.text_splitter import CharacterTextSplitter

# Import vector store for database operations
from langchain.vectorstores import Chroma

# for loading of llama gguf model
from langchain.llms import LlamaCpp

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE

from langchain.chains.router import MultiPromptChain
from langchain.chains import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory, VectorStoreRetrieverMemory
from langchain.chains import ConversationalRetrievalChain

from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler



def sanitize_collection_name(email):
    # Replace invalid characters with an underscore
    sanitized = re.sub(r'[^a-zA-Z0-9_-]', '_', email)
    # Ensure the name is within the length limits
    if len(sanitized) > 63:
        # Hashing the name to ensure uniqueness and length constraint
        hash_suffix = hashlib.sha256(email.encode()).hexdigest()[:8]
        sanitized = sanitized[:55] + "_" + hash_suffix
    # Ensure it starts and ends with an alphanumeric character
    if not re.match(r'^[a-zA-Z0-9].*[a-zA-Z0-9]$', sanitized):
        sanitized = "a" + sanitized + "1"
    return sanitized


# Modify vectordb initialization to be dynamic based on user_id
def get_vectordb_for_user(user_collection_name):
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    vectordb = Chroma(
        collection_name=user_collection_name,
        embedding_function=HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'),
        persist_directory=f"{CHROMADB_LOC}" # Optional: Separate directory for each user's data
    )
    return vectordb


def pdf_to_vec(filename, user_collection_name):
    
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    document = []
    loader = PyPDFLoader(filename)
    document.extend(loader.load()) #which library is this from?

    # Initialize HuggingFaceEmbeddings with the 'sentence-transformers/all-MiniLM-L6-v2' model for generating text embeddings
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

    # Initialize a CharacterTextSplitter to split the loaded documents into smaller chunks
    document_splitter = CharacterTextSplitter(separator='\n', chunk_size=500, chunk_overlap=100)

    # Use the splitter to divide the 'document' content into manageable chunks
    document_chunks = document_splitter.split_documents(document) #which library is this from?

    # Create a Chroma vector database from the document chunks with the specified embeddings, and set a directory for persistence
    vectordb = Chroma.from_documents(document_chunks, embedding=embeddings, collection_name=user_collection_name, persist_directory=CHROMADB_LOC) ## change to GUI path

    # Persist the created vector database to disk in the specified directory
    vectordb.persist() #this is mandatory?

    return(vectordb)
    #return collection  # Return the collection as the asset

class LlamaModelSingleton:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            print('Loading LLM model...')
            cls._instance = super(LlamaModelSingleton, cls).__new__(cls)
            
            # Model loading logic
            model_path = os.getenv("MODEL_PATH")
            cls._instance.llm = LlamaCpp(
                #streaming = True,
                model_path=model_path,
                n_gpu_layers=-1,
                n_batch=512,
                temperature=0.1,
                top_p=1,
                #verbose=False,
                #callback_manager=callback_manager,
                max_tokens=2000,
            )
            print(f'Model loaded from {model_path}')
        return cls._instance.llm


def load_llm():
    return LlamaModelSingleton()



#step 5, to instantiate once to create default_chain,router_chain,destination_chains into chain and set vectordb. so will not re-create per prompt
def default_chain(llm, user_collection_name):
    # Get Chromadb location
    CHROMADB_LOC = os.getenv('CHROMADB_LOC')

    vectordb = get_vectordb_for_user(user_collection_name)  # Use the dynamic vectordb based on user_id
    sum_template = """
    As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

    our role entails:

    Providing Detailed Explanations: Deliver comprehensive answers to these questions, elucidating the underlying technical principles.
    Assisting in Exam Preparation: Support educators in formulating sophisticated exam and quiz questions, including MCQs, accompanied by thorough explanations.
    Summarizing Course Material: Distill key information from course materials, articulating complex ideas within the context of advanced machine learning practices.

    Objective: to summarize and explain the key points.
    Here the question:
    {input}"""

    mcq_template = """
    As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

    our role entails:
    Crafting Insightful Questions: Develop thought-provoking questions that explore the intricacies of machine learning topics.
    Generating MCQs: Create MCQs for each machine learning topic, comprising a question, four choices (A-D), and the correct answer, along with a rationale explaining the answer.

    Objective: to create multiple choice question in this format
    [question:
    options A:
    options B:
    options C:
    options D:
    correct_answer:
    explanation:]

    Here the question:
    {input}"""

    prompt_infos = [
        {
            "name": "SUMMARIZE",
            "description": "Good for summarizing and explaination ",
            "prompt_template": sum_template,
        },
        {
            "name": "MCQ",
            "description": "Good for creating multiple choices questions",
            "prompt_template": mcq_template,
        },
    ]

    destination_chains = {}

    for p_info in prompt_infos:
        name = p_info["name"]
        prompt_template = p_info["prompt_template"]
        #vectordb=p_info["vector"]
        prompt = PromptTemplate(template=prompt_template, input_variables=["input"])
        embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2') #new
        vectordb= Chroma(persist_directory = CHROMADB_LOC, embedding_function = embeddings) #new
        retriever = vectordb.as_retriever()#new
        memory = VectorStoreRetrieverMemory(retriever=retriever) #new

        chain = LLMChain(llm=llm, prompt=prompt, verbose=True, memory=memory) #new memory=memory
        destination_chains[name] = chain

    #default_chain = ConversationChain(llm=llm, output_key="text")
    #memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

    default_chain = ConversationalRetrievalChain.from_llm(llm=llm,
                                                retriever=vectordb.as_retriever(search_kwargs={'k': 3}),
                                                verbose=True, output_key="text" )

    destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
    destinations_str = "\n".join(destinations)
    router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
    router_prompt = PromptTemplate(
        template=router_template,
        input_variables=["input"],
        output_parser=RouterOutputParser(),
    )
    router_chain = LLMRouterChain.from_llm(llm, router_prompt)

    return default_chain,router_chain,destination_chains

# Adjust llm_infer to accept user_id and use it for user-specific processing
def llm_infer(user_collection_name, prompt):
    
    llm = load_llm()  # load_llm is singleton for entire system
    
    vectordb = get_vectordb_for_user(user_collection_name) # Vector collection for each us. 
    
    default_chain, router_chain, destination_chains = get_or_create_chain(user_collection_name, llm)  # Now user-specific

    chain = MultiPromptChain(
        router_chain=router_chain,
        destination_chains=destination_chains,
        default_chain=default_chain,
        #memory=ConversationBufferMemory(k=2), # memory_key='chat_history', return_messages=True
        verbose=True,
    )
    response = chain.run(prompt)

    return response

# Assuming a simplified caching mechanism for demonstration
chain_cache = {}

def get_or_create_chain(user_collection_name, llm):
    if 'default_chain' in chain_cache and 'router_chain' in chain_cache:
        default_chain = chain_cache['default_chain']
        router_chain = chain_cache['router_chain']
        destination_chains = chain_cache['destination_chains']
    else:
        vectordb = get_vectordb_for_user(user_collection_name)  # User-specific vector database
        sum_template = """
        As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

        our role entails:

        Providing Detailed Explanations: Deliver comprehensive answers to these questions, elucidating the underlying technical principles.
        Assisting in Exam Preparation: Support educators in formulating sophisticated exam and quiz questions, including MCQs, accompanied by thorough explanations.
        Summarizing Course Material: Distill key information from course materials, articulating complex ideas within the context of advanced machine learning practices.

        Objective: to summarize and explain the key points.
        Here the question:
        {input}"""

        mcq_template = """
        As a machine learning education specialist, our expertise is pivotal in deepening the comprehension of complex machine learning concepts for both educators and students.

        our role entails:
        Crafting Insightful Questions: Develop thought-provoking questions that explore the intricacies of machine learning topics.
        Generating MCQs: Create MCQs for each machine learning topic, comprising a question, four choices (A-D), and the correct answer, along with a rationale explaining the answer.

        Objective: to create multiple choice question in this format
        [question:
        options A:
        options B:
        options C:
        options D:
        correct_answer:
        explanation:]

        Here the question:
        {input}"""

        prompt_infos = [
            {
                "name": "SUMMARIZE",
                "description": "Good for summarizing and explaination ",
                "prompt_template": sum_template,
            },
            {
                "name": "MCQ",
                "description": "Good for creating multiple choices questions",
                "prompt_template": mcq_template,
            },
        ]

        destination_chains = {}

        for p_info in prompt_infos:
            name = p_info["name"]
            prompt_template = p_info["prompt_template"]
            prompt = PromptTemplate(template=prompt_template, input_variables=["input"])
            chain = LLMChain(llm=llm, prompt=prompt)
            destination_chains[name] = chain
        #default_chain = ConversationChain(llm=llm, output_key="text")
        #memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

        default_chain = ConversationalRetrievalChain.from_llm(llm=llm,
                                                    retriever=vectordb.as_retriever(search_kwargs={'k': 3}),
                                                    verbose=True, output_key="text" )

        destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
        destinations_str = "\n".join(destinations)
        router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
        router_prompt = PromptTemplate(
            template=router_template,
            input_variables=["input"],
            output_parser=RouterOutputParser(),
        )
        router_chain = LLMRouterChain.from_llm(llm, router_prompt)
#
        chain_cache['default_chain'] = default_chain
        chain_cache['router_chain'] = router_chain
        chain_cache['destination_chains'] = destination_chains
    
    # Here we can adapt the chains if needed based on the user_id, for example, by adjusting the vectordb retriever
    # This is where user-specific adaptations occur

    return default_chain, router_chain, destination_chains