import os
import tempfile

from chainlit.types import AskFileResponse
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain.storage import LocalFileStore
from langchain_qdrant import QdrantVectorStore
from langchain.embeddings import CacheBackedEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.globals import set_llm_cache
from langchain_openai import ChatOpenAI
from langchain_core.caches import InMemoryCache
from langchain_core.runnables.passthrough import RunnablePassthrough
from uuid import uuid4
from utilities.prompts import get_system_template, get_user_template


def load_file(file: AskFileResponse, chunk_size=1000, chunk_overlap=100):
    import tempfile
    with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tempfile:
        with open(tempfile.name, "wb") as f:
            f.write(file.content)

    Loader = PyMuPDFLoader

    loader = Loader(tempfile.name)
    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = text_splitter.split_documents(documents)
    for i, doc in enumerate(docs):
        doc.metadata["source"] = f"source_{i}"
    return docs
    

def process_embeddings(docs):
    core_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

    collection_name = f"pdf_to_parse_{uuid4()}"
    client = QdrantClient(":memory:")
    client.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
    )
    # Adding cache!
    store = LocalFileStore("./cache/")
    cached_embedder = CacheBackedEmbeddings.from_bytes_store(
        core_embeddings, store, namespace=core_embeddings.model
    )
    # Typical QDrant Vector Store Set-up
    vectorstore = QdrantVectorStore(
        client=client,
        collection_name=collection_name,
        embedding=cached_embedder)
    vectorstore.add_documents(docs)
    retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
    return retriever


def prepare_rag_chain(retriever, prompt_cache="yes"):
    print(prompt_cache)
    system_template = get_system_template()
    user_template = get_user_template()

    chat_prompt = ChatPromptTemplate.from_messages([
        ("system", system_template),
        ("human", user_template)
    ])

    chat_model = ChatOpenAI(model="gpt-4o-mini")

    if prompt_cache == "yes":
        set_llm_cache(InMemoryCache())

    from operator import itemgetter

    rag_qa_chain = (
            {"context": itemgetter("question") | retriever, "question": itemgetter("question"), "language": itemgetter("language")}
            | RunnablePassthrough.assign(context=itemgetter("context"), language=itemgetter("language"))
            | chat_prompt | chat_model
        )
    return rag_qa_chain 

def process_file(file, prompt_cache):
    docs = load_file(file)
    retriever = process_embeddings(docs)
    rag_chain = prepare_rag_chain(retriever, prompt_cache)
    return {"chain": rag_chain, "retriever": retriever}