from openai import AsyncOpenAI  # importing openai for API usage
import chainlit as cl  # importing chainlit for our app
from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
from dotenv import load_dotenv
import asyncio
import datetime

from aimakerspace.text_utils import TextFileLoader, CharacterTextSplitter
from aimakerspace.vectordatabase import VectorDatabase
from aimakerspace.openai_utils.prompts import (
    UserRolePrompt,
    SystemRolePrompt,
)

load_dotenv()

RAQA_PROMPT_TEMPLATE = """
Use the provided context to answer the user's query. 

You may not answer the user's query unless there is specific context in the following text.

If you do not know the answer, or cannot answer, please respond with "I don't know".

Context:
{context}
"""
raqa_prompt = SystemRolePrompt(RAQA_PROMPT_TEMPLATE)

USER_PROMPT_TEMPLATE = """
User Query:
{user_query}
"""
user_prompt = UserRolePrompt(USER_PROMPT_TEMPLATE)

vector_db = None

@cl.on_chat_start
async def start_chat():
    #Create the Vector Database for King Lear
    global vector_db
    text_loader = TextFileLoader("data/KingLear.txt")
    documents = text_loader.load_documents()
    text_splitter = CharacterTextSplitter()
    split_documents = text_splitter.split_texts(documents)
    vector_db = VectorDatabase()
    vector_db = asyncio.run(vector_db.abuild_from_list(split_documents))
    
    settings = {
        "model": "gpt-3.5-turbo",
        "temperature": 0,
        "max_tokens": 500,
        "top_p": 1,
        "frequency_penalty": 0,
        "presence_penalty": 0,
    }
    cl.user_session.set("settings", settings)

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

    client = AsyncOpenAI()

    context_list = vector_db.search_by_text(message.content, k=4)
    context_prompt = ""
    for context in context_list:
        context_prompt += context[0] + "\n"
    formatted_system_prompt = raqa_prompt.create_message(context=context_prompt)
    formatted_user_prompt = user_prompt.create_message(user_query=message.content)
    print(formatted_system_prompt)
    print(formatted_user_prompt)
    prompt = Prompt(
        provider=ChatOpenAI.id,
        messages=[
            PromptMessage(
                role="system",
                template=RAQA_PROMPT_TEMPLATE,
                formatted=formatted_system_prompt['content'],
            ),
            PromptMessage(
                role="user",
                template=USER_PROMPT_TEMPLATE,
                formatted=formatted_user_prompt['content'],
            ),
        ],
        inputs={"context": context_prompt,
            "user_query": message.content},
        settings=settings,
    )

    msg = cl.Message(content="")
    
    async for stream_resp in await client.chat.completions.create(
        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
    ):
        token = stream_resp.choices[0].delta.content
        if not token:
            token = ""
        await msg.stream_token(token)
        
    # Update the prompt object with the completion
    prompt.completion = msg.content
    msg.prompt = prompt

    # Send and close the message stream
    await msg.send()