from typing import List, Tuple, Dict, Generator
from langchain.llms import OpenAI
import gradio as gr


model_name = "gpt-3.5-turbo"
LLM = OpenAI(model_name=model_name, temperature=0.9)


def create_history_messages(history: List[Tuple[str, str]]) -> List[dict]:
    history_messages = [{"role": "user", "content": m[0]} for m in history]
    history_messages.extend([{"role": "assistant", "content": m[1]} for m in history])
    return history_messages


def create_formatted_history(history_messages: List[dict]) -> List[Tuple[str, str]]:
    formatted_history = []
    user_messages = []
    assistant_messages = []

    for message in history_messages:
        if message["role"] == "user":
            user_messages.append(message["content"])
        elif message["role"] == "assistant":
            assistant_messages.append(message["content"])

        if user_messages and assistant_messages:
            formatted_history.append(
                ("".join(user_messages), "".join(assistant_messages))
            )
            user_messages = []
            assistant_messages = []

    # append any remaining messages
    if user_messages:
        formatted_history.append(("".join(user_messages), None))
    elif assistant_messages:
        formatted_history.append((None, "".join(assistant_messages)))

    return formatted_history


history_messages = [{"role": "system", "content": "A helpful assistant."}]


def chat(
    message: str, history=[], client=LLM.client
) -> Generator[Tuple[List[Tuple[str, str]], List[Tuple[str, str]]], None, None]:
    global history_messages
    history_messages.append({"role": "user", "content": message})

    # We have no content for the assistant's response yet but we will update this:
    history_messages.append({"role": "assistant", "content": ""})

    response_message = ""

    chat_generator = client.create(
        messages=history_messages, stream=True, model=model_name
    )

    for chunk in chat_generator:
        if "choices" in chunk:
            for choice in chunk["choices"]:
                if "delta" in choice and "content" in choice["delta"]:
                    new_token = choice["delta"]["content"]
                    # Add the latest token:
                    response_message += new_token
                    # Update the assistant's response in our model:
                    history_messages[-1]["content"] = response_message

                if "finish_reason" in choice and choice["finish_reason"] == "stop":
                    break
        formatted_history = create_formatted_history(history_messages)
        yield formatted_history, formatted_history


# Define the settings for the Gradio ChatBot interface
chatbot = gr.Chatbot(label="Chat").style(color_map=("yellow", "purple"))

# Define the main interface with input components for the API key, message, and state
iface = gr.Interface(
    fn=chat,
    inputs=[
        gr.Textbox(placeholder="Hello! How are you? etc.", label="Message"),
        "state",
    ],
    outputs=[chatbot, "state"],
    allow_flagging="never",  # Hide the "flag" button
)

# Launch the Gradio interface
iface.queue().launch()


# TODO: figure out how to chain with the Chat API. Or just wait for others to do it.,, The API is only a few days old, afterall.
history_messages = [
    {
        "role": "system",
        "content": "A helpful assistant that performs web searches.\nA helpful assistant that performs web searches.\nAssistant can say these verbs:\n\n- `Thought: I need to search for 'search term'`\n\n- `Thought: I need to read the URL: 'URL'`\n\n- `Thought: I have the information I need.`\n\n- `Say: 'message'`\nThe date is 3rd March 2023.",
    },
    {"role": "user", "content": "Who is the President of the USA?"},
    {
        "role": "assistant",
        "content": "Thought: I need to search for the latest infomation.",
    },
    {"role": "assistant", "content": "search('USA president Wikipedia')"},
    {
        "role": "system",
        "content": "Search results for 'USA president Wikipedia'.\n\nhttps://en.wikipedia.org/wiki/President_of_the_United_States\nPresident of the United States - Wikipedia\nThe president of the United States ( POTUS) [A] is the head of state and head of government of the United States of America. The president directs the executive branch of the federal government and is the commander-in-chief of the United States Armed Forces . The power of the presidency has grown substantially [11] since the first president ...",
    },
    {
        "role": "assistant",
        "content": "load_url_text('https://en.wikipedia.org/wiki/President_of_the_United_States')",
    },
    {
        "role": "system",
        "content": "President of the United States\nPresident of the\nUnited States of America\nSeal of the President of the United States.svg\n\nPresidential seal\nFlag of the President of the United States.svg\n\nPresidential flag\n\nJoe Biden presidential portrait (cropped).jpg\n\nIncumbent\nJoe Biden\nsince January 20, 2021\n\n    Executive branch of the U.S. government\n    Executive Office of the President\n\nStyle\t\n\n",
    },
    {"role": "assistant", "content": "Thought: I have the information I need."},
    {
        "role": "assistant",
        "content": "Say: The current President of the USA is Joe Biden.",
    },
]