"""
Written in 2024 by retoor@molodetz.nl.

MIT license. Enjoy!

The purpose of this file is to be a native part of your application 
instead of yet another library. It's just not worth making a library, 
especially not another one. Just modify and use it!

The docstrings of all methods contain tips and important facts.
This document contains all URLs for all services that you need.

You'll need:
 - An OpenAI account.
 - A named project in the OpenAI dashboard.
 - A requested API key and an assistant created.

URLs to all these services are described in the class for convenience.
They can be hard to find initially.

The API keys described in this document are fake but are in the correct format for educational purposes.

How to start:
 - sudo apt install python3.12-venv python3-pip -y
 - python3 -m venv .venv
 - . .venv/bin/activate
 - pip install openai

"""

# AGAIN, NOT REAL DATA, ONLY LOOKS LIKE IT FOR EDUCATIONAL PURPOSES. 
# Not required to use the Agent class. The Agent class accepts api_key as a parameter.
API_KEY = "sk-proj-V1Jc3my22xSvtfZ3dxXNHgLWZIhEopmJVIMlcNrft_q-7p8dDT_-AQCE8wo9cKpO3v05egDm7CT3BlbkFjN21maiSZqS4oz8FSGiblOeKMH2i6BzIGdQWMcVbKHnRqWy0KiSwKQywJ7XEf792UgGFtwLtxkA"
# Not required to use the Agent class. The Agent class accepts assistant_id as a parameter.
ASSISTANT_ID = "asst_NgncvKEN8CTf642RE8a4PgAp"


import asyncio
import functools
from collections.abc import Generator
from typing import Optional

from openai import OpenAI


class Agent:
    """
    This class represents a single user session with its own memory.

    The messages property of this class is a list containing the full chat history about
    what the user said and what the assistant (agent) said. This can be used in the future to continue
    where you left off. The format is described in the docs of the __init__ function below.

    Introduction to API usage if you want to extend this class:
    https://platform.openai.com/docs/api-reference/introduction
    """

    def __init__(
        self, api_key: str, assistant_id: int, messages: Optional[list] = None
    ):
        """
        You can find and create API keys here:
        https://platform.openai.com/api-keys

        You can find the assistant_id (agent_id) here. It is the ID that starts with 'asst_', not your custom name:
        https://platform.openai.com/assistants/

        Messages are optional and should be in this format to keep a message history that you can later use again:
        [
            {"role": "user", "message": "What is choking the chicken?"},
            {"role": "assistant", "message": "Lucky for the cock."}
        ]
        """

        self.assistant_id = assistant_id
        self.api_key = api_key
        self.client = OpenAI(api_key=self.api_key)
        self.messages = messages or []
        self.thread = self.client.beta.threads.create(messages=self.messages)

    async def dalle2(
        self, prompt: str, width: Optional[int] = 512, height: Optional[int] = 512
    ) -> dict:
        """
        In my opinion, DALL·E 2 produces unusual results.
        Sizes: 256x256, 512x512, or 1024x1024.
        """
        result = self.client.images.generate(
            model="dall-e-2", prompt=prompt, n=1, size=f"{width}x{height}"
        )
        return result

    @property
    async def models(self):
        """
        List models in dict format. That's more convenient than the original
        list method because this can be directly converted to JSON to be used
        in your frontend or API. This is not the original result, which is a
        custom list with unserializable models.
        """
        return [
            {
                "id": model.id,
                "owned_by": model.owned_by,
                "object": model.object,
                "created": model.created,
            }
            for model in self.client.models.list()
        ]

    async def dalle3(
        self, prompt: str, height: Optional[int] = 1024, width: Optional[int] = 1024
    ) -> dict:
        """
        Sadly, only large sizes are allowed. It's more expensive.
        Sizes: 1024x1024, 1792x1024, or 1024x1792.
        """
        result = self.client.images.generate(
            model="dall-e-3", prompt=prompt, n=1, size=f"{width}x{height}"
        )
        print(result)
        return result

    async def chat(
        self, message: str, interval: Optional[float] = 0.2
    ) -> Generator[None, None, str]:
        """
        Chat with the agent. It yields at the given interval to inform the caller it's still busy, so you can
        update the user with a live status. It doesn't hang. You can use this fully asynchronously with other
        instances of this class.

        This function also updates the self.messages list with chat history for later use.
        """
        message_object = {"role": "user", "content": message}
        self.messages.append(message_object)
        self.client.beta.threads.messages.create(
            self.thread.id,
            role=message_object["role"],
            content=message_object["content"],
        )
        run = self.client.beta.threads.runs.create(
            thread_id=self.thread.id, assistant_id=self.assistant_id
        )

        while run.status != "completed":
            run = self.client.beta.threads.runs.retrieve(
                thread_id=self.thread.id, run_id=run.id
            )
            yield None
            await asyncio.sleep(interval)

        response_messages = self.client.beta.threads.messages.list(
            thread_id=self.thread.id
        ).data
        last_message = response_messages[0]
        self.messages.append({"role": "assistant", "content": last_message})
        yield last_message

    async def chatp(self, message: str) -> str:
        """
        Just like the regular chat function but with progress indication and returns a string directly.
        This is handy for interactive usage or for a process log.
        """
        asyncio.get_event_loop()
        print("Processing", end="")
        async for message in self.chat(message):
            if not message:
                print(".", end="", flush=True)
                continue
            print("")
            break
        return message

    async def read_line(self, ps: Optional[str] = "> "):
        """
        Non-blocking read_line.
        Blocking read_line can break WebSocket connections.
        That's why.
        """
        loop = asyncio.get_event_loop()
        patched_input = functools.partial(input, ps)
        return await loop.run_in_executor(None, patched_input)

    async def cli(self):
        """
        Interactive client. Can be used in a terminal by the user or a different process.
        The bottom newline is so that a process can check for '\n\n' to determine if the response has ended
        and there's nothing left to wait for, allowing the process to send the next prompt if the '>' shows.
        """
        while True:
            try:
                message = await self.read_line("> ")
                if not message.strip():
                    continue
                response = await self.chatp(message)
                print(response.content[0].text.value)
                print("")
            except KeyboardInterrupt:
                print("Exiting...")
                break

async def main():
    """
    Example main function. The keys here are not real but look exactly like 
    the real ones for example purposes so you can verify your key is in the 
    correct format.
    """
    agent = Agent(api_key=API_KEY, assistant_id=ASSISTANT_ID)
    
    # Generate an image. Use DALL·E 3, as DALL·E 2 is almost unusable. For image sizes, look at the class method docstring.
    list_containing_dicts_with_url_to_images = await agent.dalle3("Make a photo-realistic image of a Rust developer")
    
    # Run interactive chat
    await agent.cli()


if __name__ == "__main__":
    # Only executed by direct execution of the script, not when imported.
    asyncio.run(main())