from dotenv import load_dotenv

load_dotenv()

import logging
import os
from time import perf_counter

import chainlit as cl
from langchain.cache import InMemoryCache
from langchain.embeddings import CacheBackedEmbeddings
from langchain.globals import set_llm_cache
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from langchain.schema import SystemMessage
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
from langchain.storage import LocalFileStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai.chat_models import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from pinecone import Index

from src.loaders import get_docs

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
os.environ["LANGCHAIN_WANDB_TRACING"] = "true"

try:
    set_llm_cache(InMemoryCache())  # Use in-memory cache for LLM

    index = Index(
        api_key=os.environ["PINECONE_API_KEY"],
        index_name=os.environ["PINECONE_INDEX"],
        host=os.environ["PINECONE_HOST"],
    )

    store = LocalFileStore("./cache/")
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    embeddings = CacheBackedEmbeddings.from_bytes_store(
        underlying_embeddings=embeddings,
        document_embedding_cache=store,
        namespace=embeddings.model,
    )

    vectorstore = PineconeVectorStore(
        index=index, embedding=embeddings, text_key="text"
    )

    system_message = SystemMessage(content="You are a helpful assistant.")
    human_template = """Based on the following context generate an response for the user query. If the response is not available say I dont know.
    Context: {context}

    Query: {query}

    Answer:"""
    human_message = HumanMessagePromptTemplate.from_template(template=human_template)
    llm = ChatOpenAI(model="gpt-3.5-turbo-0125", streaming=True)
    prompt = ChatPromptTemplate.from_messages(messages=[system_message, human_message])
    parser = StrOutputParser()
    prompt_chain = prompt | llm | parser
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=250, chunk_overlap=25, length_function=len
    )
    logger.info("Initialization completed successfully.")
except Exception as e:
    logger.exception("Failed during initialization: %s", str(e))


@cl.on_chat_start
async def start_chat():
    files = await cl.AskFileMessage(
        content="Please upload a PDF file to answer questions from.",
        accept=["application/pdf"],
        max_size_mb=100,
        max_files=10,
        timeout=60 * 60 * 24 * 7 * 365,
    ).send()

    out = cl.Message(content="")
    await out.send()

    paths = [file.path for file in files]
    logger.info(files[0])
    logger.info("Preparing docs: %s", paths)
    start = perf_counter()
    splitted_docs = get_docs(files=files, splitter=splitter)
    end = perf_counter()
    logger.info("Preparing docs took %s seconds.", end - start)

    retriever = vectorstore.as_retriever(
        search_kwargs={"filter": {"source": {"$in": paths}}}
    )

    logger.info("Adding documents to vector store retriever.")
    start = perf_counter()
    await retriever.aadd_documents(splitted_docs)
    end = perf_counter()
    logger.info("Adding documents took %s seconds.", end - start)

    cl.user_session.set("retriever", retriever)
    out.content = f"{len(files)} file(s) loaded! You can now ask questions"
    await out.update()
    logger.info("Files loaded and retriever updated.")


@cl.on_message
async def main(message: cl.Message):
    retriever = cl.user_session.get("retriever")
    retriever_chain = RunnableParallel(
        {"context": retriever, "query": RunnablePassthrough()}
    )

    out = cl.Message(content="")
    await out.send()

    chain = retriever_chain | prompt_chain
    stream = chain.astream(message.content)

    async for chunk in stream:
        await out.stream_token(chunk)

    await out.update()