import os
from operator import itemgetter

import chainlit as cl
from langchain import hub
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import AIMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import ConfigurableFieldSpec, RunnablePassthrough
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from lib.in_memory_history import InMemoryHistory
from lib.indexing import (create_search_engine, format_docs,
                          load_and_split_documents)
from lib.prompts import EXAMPLE_PROMPT, PROMPT, WELCOME_MESSAGE

llm_model = "gpt-3.5-turbo"
temperature=0.0

store = {}

# Fetch message history
def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHistory:
    if (user_id, conversation_id) not in store:
        store[(user_id, conversation_id)] = InMemoryHistory()
    return store[(user_id, conversation_id)]


history = get_session_history("1", "1")
history.add_message(AIMessage(content="hello"))
print(store)


@cl.on_chat_start
async def init():
    """This function is run at every chat session starts to ask user for file,
    index it, and build the RAG chain.

    Raises:
        SystemError: yolo
    """

    await cl.Avatar(
        name="Chatbot",
        url="https://www.aicatalystpartners.com/hubfs/favicon/ms-icon-310x310.png",
    ).send()

    files = None

    # Asking user to to upload a PDF to chat with
    while files == None:
        files = await cl.AskFileMessage(
            content=WELCOME_MESSAGE,
            accept=["application/pdf"],
            max_size_mb=24,
            timeout=180,
        ).send()

    file = files[0]

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

    docs = load_and_split_documents(file=file)
    cl.user_session.set("docs", docs)

    msg.content = f"The file '{file.name}' has been processed. Loading index ..."
    await msg.update()

    # Index documents into search engine
    embeddings = OpenAIEmbeddings(model="text-embedding-ada-002") # we can define a different embedding model in future
    try:
        search_engine = await cl.make_async(create_search_engine)(
            docs=docs, embeddings=embeddings
        )
    except Exception as e:
        await cl.Message(content=f"Error: {e}").send()
        raise SystemError
    
    # Prompt - using an external prompt template
    prompt = hub.pull("rlm/rag-prompt")
    
    # LLM model
    llm = ChatOpenAI(temperature=temperature, model=llm_model, streaming=True)

    # Output parser
    output_parser = StrOutputParser()

    # RAG Chain
    context = itemgetter("question") | search_engine.as_retriever(search_kwargs={'k': 4}, max_tokens_limit=4097) | format_docs
    first_step = RunnablePassthrough.assign(context=context)
    rag_chain = first_step | prompt | llm
    
    rag_chain_with_mgs_history = RunnableWithMessageHistory(
        rag_chain, # type: ignore
        get_session_history=get_session_history,
        input_messages_key="question",
        history_messages_key="history",
        history_factory_config=[
            ConfigurableFieldSpec(
                id="user_id",
                annotation=str,
                name="User ID",
                description="Unique identifier for the user.",
                default="",
                is_shared=True,
            ),
            ConfigurableFieldSpec(
                id="conversation_id",
                annotation=str,
                name="Conversation ID",
                description="Unique identifier for the conversation.",
                default="",
                is_shared=True,
            ),
        ],
    )  | output_parser

    # Let the user know that the RAG QA system is ready
    msg.content = f"The file '{file.name}' has been processed and the index has been loaded. You can now ask questions!"
    await msg.update()

    cl.user_session.set("runnable", rag_chain_with_mgs_history)


@cl.on_message
async def on_message(message: cl.Message):
    """This function is invoked whenever we receive a Chainlit message.

    Args:
        message (cl.Message): user input
    """

    runnable_chain = cl.user_session.get("runnable")  # type: ignore # type: 
    cb = cl.AsyncLangchainCallbackHandler(stream_final_answer=True)

    response = await runnable_chain.ainvoke( # type: ignore # type: RunnableWithMessageHistory
        {"question": message.content},
        config = {
            "configurable": {"user_id": "user_id", "conversation_id": "conversation_id"},
            "callbacks": [cb]
        },
    )
    
    await cl.Message(content=response).send()

    msg = cl.Message(content="")