import os
import numpy as np
from typing import List
from chainlit.types import AskFileResponse
from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader
from aimakerspace.openai_utils.prompts import (
    UserRolePrompt,
    SystemRolePrompt,
    AssistantRolePrompt,
)
from aimakerspace.openai_utils.embedding import EmbeddingModel
from aimakerspace.vectordatabase import VectorDatabase
from aimakerspace.openai_utils.chatmodel import ChatOpenAI
import chainlit as cl
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
from chainlit.input_widget import Select
from qdrant_client.models import PointStruct
#Qdrant client
client = None

#System Chat Prompt
system_template = """\
Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
system_role_prompt = SystemRolePrompt(system_template)

#User Prompt for chat
user_prompt_template = """\
Context:
{context}

Question:
{question}
"""
user_role_prompt = UserRolePrompt(user_prompt_template)

#Categorization of VectorDatabase
system_template_db = """\
You are an expert in categorization. Given the last user response determine if he or she wants to use the Qdrant database. If yes return the output single word 'QDrant' without any other phrases. If no return the only the word 'AI Makerspace'.

"""
system_role_prompt_db = SystemRolePrompt(system_template_db)

user_prompt_template_db = "User Input:\n{user_input}"
user_role_prompt_db = UserRolePrompt(user_prompt_template)

class RetrievalAugmentedQAPipeline:
    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
        self.llm = llm
        self.vector_db_retriever = vector_db_retriever

    async def arun_pipeline(self, user_query: str):
        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)

        context_prompt = ""
        for context in context_list:
            context_prompt += context[0] + "\n"

        formatted_system_prompt = system_role_prompt.create_message()

        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)

        async def generate_response():
            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):
                yield chunk

        return {"response": generate_response(), "context": context_list}

class RetrievalAugmentedQAPipelineQdrant:
    def __init__(self, llm: ChatOpenAI(), vector_db_retriever) -> None:
        self.llm = llm
        self.vector_db_retriever = vector_db_retriever
        self.embedding_model = EmbeddingModel()

    async def arun_pipeline(self, user_query: str):
        query_vector = self.embedding_model.get_embedding(user_query)
        context_list = self.vector_db_retriever.search(
            collection_name="my_collection",
            query_vector=query_vector, 
            limit=4
        )

        context_prompt = ""
        for context in context_list:
            context_prompt += context.payload['text'] + "\n"

        formatted_system_prompt = system_role_prompt.create_message()

        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)

        async def generate_response():
            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):
                yield chunk

        return {"response": generate_response(), "context": context_list}

text_splitter = CharacterTextSplitter()


def process_text_file(file: AskFileResponse):
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as temp_file:
        temp_file_path = temp_file.name

    with open(temp_file_path, "wb") as f:
        f.write(file.content)

    text_loader = TextFileLoader(temp_file_path)
    documents = text_loader.load_documents()
    texts = text_splitter.split_texts(documents)
    return texts

def process_pdf_file(file: AskFileResponse):
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pdf") as temp_file:
        temp_file_path = temp_file.name

    with open(temp_file_path, "wb") as f:
        f.write(file.content)

    text_loader = TextFileLoader(temp_file_path)
    documents = text_loader.load_documents()
    texts = text_splitter.split_texts(documents)
    return texts

async def initialize_qdrant(text):
    client = QdrantClient(":memory:")
    if not client.collection_exists("my_collection"):
        client.create_collection(
            collection_name="my_collection",
            vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
        )
    embedding_model = EmbeddingModel()
    embeddings = await embedding_model.async_get_embeddings(text)
    i = 0
    for text, embedding in zip(text, embeddings):
        insert(text, np.array(embedding), i, client)
        i+=1
   
    return client

def insert(text, vector, idx, client):
    point= PointStruct(
                id=idx,
                vector=vector.tolist(),
                payload={"text": text}
    )
    client.upsert(
        collection_name="my_collection",
        points=[point]
    )
    
    
def choose_db(llm, user_input):
    formatted_system_prompt_db = system_role_prompt_db.create_message()
    formatted_user_prompt_db = user_role_prompt_db.create_message(question=user_input)
    return llm.run([formatted_system_prompt_db, formatted_user_prompt_db])

@cl.on_chat_start
async def on_chat_start():
    global client
    files = None

    # Wait for the user to upload a file
    while files == None :

        files = await cl.AskFileMessage(
            content="Please upload a Text or PDF file to begin!",
            accept=["text/plain", "application/pdf"],
            max_size_mb=2,
            timeout=180,
        ).send()


    file = files[0]

    msg = cl.Message(
        content=f"Processing `{file.name}`...", disable_human_feedback=True
    )
    await msg.send()

    # load the file
    if file.name.endswith('.pdf'):
        texts = process_pdf_file(file)
    else:
        texts = process_text_file(file)

    print(f"Processing {len(texts)} text chunks")
    
    chat_openai = ChatOpenAI()

    res = await cl.AskUserMessage(content="Do you want to use the QDrant vector database or AI Makerspace's?").send()
    
    if res:
        chosen_db = choose_db(chat_openai, res['content'])
        await cl.Message(
            content=f"You have chosen {chosen_db}. Please start asking questions!",
        ).send()

    # Create a chain
    retrieval_augmented_qa_pipeline = None
    if chosen_db.lower() == 'qdrant': 
       client = await initialize_qdrant(texts)
       retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipelineQdrant(
            vector_db_retriever=client,
            llm=chat_openai
        )
    else:  
        vector_db = VectorDatabase()
        vector_db = await vector_db.abuild_from_list(texts)
        retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
            vector_db_retriever=vector_db,
            llm=chat_openai
        )
    
    # Let the user know that the system is ready
    msg.content = f"Processing `{file.name}` done. You can now ask questions!"
    await msg.update()

    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)


@cl.on_message
async def main(message):
    chain = cl.user_session.get("chain")

    msg = cl.Message(content="")
    result = await chain.arun_pipeline(message.content)

    async for stream_resp in result["response"]:
        await msg.stream_token(stream_resp)

    await msg.send()