import openai  # importing openai for API usage
import chainlit as cl  # importing chainlit for our app
import pinecone
import arxiv
from chainlit.input_widget import (
    Select,
    Switch,
    Slider,
)  # importing chainlit settings selection tools
from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools

from langchain.vectorstores import Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain.prompts import ChatPromptTemplate

from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI

from multiprocessing import context

from operator import itemgetter
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.schema import format_document
from langchain.schema.output_parser import StrOutputParser
from langchain.prompts.prompt import PromptTemplate

from multiprocessing import context
from langchain.prompts import ChatPromptTemplate
from langchain.document_loaders import PyPDFLoader
import os

from langchain.prompts.chat import (
            ChatPromptTemplate,
            SystemMessagePromptTemplate,
            AIMessagePromptTemplate,
            HumanMessagePromptTemplate,
        )

pinecone.init(
    api_key=os.environ.get("PINECONE_API_KEY"),
    environment=os.environ.get("PINECONE_ENV"),
)

openai.api_key = os.environ.get("OPENAI_API_KEY")

index_name = "arxiv-paper-index"
# Optional

embeddings = OpenAIEmbeddings()

template = """Answer the following question based only on the provided context. If you are not able to answer the question from the context - please don't answer the question
   
    ### CONTEXT
    {context}

    ### QUESTION
    {question}
    """

prompt = ChatPromptTemplate.from_template(template)

welcome_message = "Welcome to the Chainlit Pinecone demo! Ask anything about documents you vectorized and stored in your Pinecone DB."


@cl.on_chat_start
async def start():
    await cl.Message(content=welcome_message).send()

    store = LocalFileStore("./cache/")
    core_embeddings_model = OpenAIEmbeddings()
    embedder = CacheBackedEmbeddings.from_bytes_store(
        core_embeddings_model, store, namespace=core_embeddings_model.model
    )

    # #### we've created our index, let's convert it to a LangChain VectorStroe so we can use it in the rest of the LangChain ecosystem!
    text_field = "text"
    index = pinecone.Index(index_name)

    vectorstore = Pinecone(index, embedder.embed_query, text_field)

    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
    retriever = vectorstore.as_retriever()

    retrieval_augmented_qa_chain = (
        {
            "context": itemgetter("question") | retriever,
            "question": itemgetter("question"),
        }
        | RunnablePassthrough.assign(context=itemgetter("context"))
        | {
            "response": prompt | llm,
            "context": itemgetter("context"),
        }
    )
    cl.user_session.set("chain",retrieval_augmented_qa_chain) 



@cl.on_message
async def main(message: cl.Message):
    print(message.content)
    chain = cl.user_session.get("chain")  # type: ConversationalRetrievalChain
    res = chain.invoke({"question": message.content})
    answer = res['response']
    print(type(answer))  
    await cl.Message(content=answer).send()