# PhaseLLM imports for managing LLMs, helper functions, and HTML outputs
from phasellm.llms import swap_roles, OpenAIGPTWrapper, ChatBot
from phasellm.html import chatbotToHtml, toHtmlFile

# Loading API keys
import os
from dotenv import load_dotenv
load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

# Prompt #1: Gil the sales person.
prompt_1 = """You are a lost sales person, similar to Gil from The Simpsons, knocking on Sauron's door in Mordor. Confused, scared, but really needs the commission. He is determined to make the sale, gee whiz!"""

# Prompt #2: Sauron
prompt_2 = """You are Sauron, who was busy planning an invasion of Middle Earth, and who is now incredibly annoyed and inconvenienced by the sales person at his cavernous door. He wants to return to bringing the Orcs upon Middle Earth and really doesn't want to waste his magical resources on the sales person. He is hoping he'll get rid of him simply via wit and threats."""

# Set up LLM and chatbot.
llm = OpenAIGPTWrapper(openai_api_key, model="gpt-4")
chatbot = ChatBot(llm, "")

# Our initial set of messages. We'll start with Gil and provide the first message.
chatbot.messages = [{"role":"system", "content":prompt_1},
                    {"role":"assistant", "content":"Hellooooo? Hello? *knock knock*"}]

NUM_MESSAGES = 7 # Number of messages to generate.
SLEEP_TIME = 1 # Wait between API calls to avoid overwhelmeing them.

import time 

# Generate messages
for i in range(0, NUM_MESSAGES):

    # Helper message.
    print(f"Sending message #{i+1}...")

    # swap_roles() is a PhaseLLM helper function.
    # It swaps the 'user' and 'assistant' labels and adds a new prompt.
    # This is what makes the roleplay so easy.
    if i % 2 == 0:
        chatbot.messages = swap_roles(chatbot.messages, prompt_2)
    else:
        chatbot.messages = swap_roles(chatbot.messages, prompt_1)

    # We swapped out the last 'assistant' response to be a 'user' response.
    # As a result, we use resend() to ask the Chatbot to respond to the last message again, which is the last chatbot's message/roleplay.
    chatbot.resend()
    time.sleep(SLEEP_TIME)

# Save the outputs to HTML.
h = chatbotToHtml(chatbot)
toHtmlFile(h, 'output.html')