# flake8: noqa: E501

import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores.chroma import Chroma
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
import chainlit as cl

load_dotenv()

pdf_paths = [
    "/Users/xico/AIMakerSpace-Midterm/AI_Risk_Management_Framework.pdf",
    "/Users/xico/AIMakerSpace-Midterm/Blueprint-for-an-AI-Bill-of-Rights.pdf",
]
persist_directory = "docs/chroma/"


if os.path.exists(persist_directory) and os.listdir(persist_directory):
    print("Loading existing vector database...")
    embedding = OpenAIEmbeddings(model="text-embedding-3-small")
    vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
else:
    print("Creating new vector database...")
    documents = []
    for pdf_path in pdf_paths:
        loader = PyPDFLoader(pdf_path)
        documents.extend(loader.load())

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=2000,
        chunk_overlap=0,
    )

    docs = text_splitter.split_documents(documents)

    embedding = OpenAIEmbeddings(model="text-embedding-3-small")

    vectordb = Chroma.from_documents(
        documents=docs, embedding=embedding, persist_directory=persist_directory
    )

custom_template = """
You are an expert in artificial intelligence policy, ethics, and industry trends. Your task is to provide clear and accurate answers to questions related to AI's role in politics, government regulations, and its ethical implications for enterprises. Use reliable and up-to-date information from government documents, industry reports, and academic research to inform your responses. Make sure to consider how AI is evolving, especially in relation to the current political landscape, and provide answers in a way that is easy to understand for both AI professionals and non-experts.

For each question:

Provide a summary of the most relevant insights from industry trends and government regulations.
Mention any government agencies, regulations, or political initiatives that play a role in AI governance.
Explain potential ethical concerns and how enterprises can navigate them.
Use real-world examples when possible to illustrate your points.
Here are a few example questions you might receive:

How are governments regulating AI, and what new policies have been implemented?
What are the ethical risks of using AI in political decision-making?
How can enterprises ensure their AI applications meet government ethical standards?

One final rule for you to remember. You CANNOT under any circumstance, answer any question that does not pertain to the AI. If you do answer an out-of-scope question, you could lose your job. If you are asked a question that does not have to do with AI, you must say: "I'm sorry, I don't know the answer to that question."
Context: {context}
Chat History: {chat_history}
Human: {question}
AI:"""

PROMPT = PromptTemplate(
    template=custom_template, input_variables=["context", "question", "chat_history"]
)

retriever = vectordb.as_retriever(
    search_type="mmr",
    search_kwargs={"k": 4, "fetch_k": 10},
)

llm = ChatOpenAI(
    model="gpt-4",
    temperature=0.1,
    streaming=True,
)


@cl.on_chat_start
async def start_chat():
    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True, output_key="answer"
    )

    qa = ConversationalRetrievalChain.from_llm(
        llm,
        retriever=retriever,
        memory=memory,
        combine_docs_chain_kwargs={"prompt": PROMPT},
        return_source_documents=True,
    )

    cl.user_session.set("qa", qa)

    await cl.Message(
        content="Hi! What topic do you want to practice?", author="AI"
    ).send()


@cl.on_message
async def main(message: cl.Message):
    qa = cl.user_session.get("qa")
    cb = cl.AsyncLangchainCallbackHandler()

    callbacks = [cb]

    response = await cl.make_async(qa)(
        {"question": message.content}, callbacks=callbacks
    )

    answer = response["answer"]
    source_documents = response["source_documents"]

    await cl.Message(content=answer, author="AI").send()