from openai.types.chat import (
    ChatCompletionAssistantMessageParam,
    ChatCompletionContentPartImageParam,
    ChatCompletionContentPartTextParam,
    ChatCompletionMessageParam,
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
)
from openai.types.chat.chat_completion_content_part_image_param import ImageURL


class Context(list[ChatCompletionMessageParam]):
    def append_system_content(self, content: str) -> None:
        """Append a system message to the context.

        Args:
            content (str): The content of the system message.
        """
        self.append(ChatCompletionSystemMessageParam(role='system', content=content))

    def append_assistant_content(self, thinking: str, action: str, result: str) -> None:
        """Append an assistant message to the context.

        Args:
            thinking (str): The assistant's thought process.
            action (str): The action the assistant wants to take.
            result (str): The result of the action.
        """
        self.append(
            ChatCompletionAssistantMessageParam(
                role='assistant', content=f'<think>{thinking}</think><answer>{action}</answer><result>{result}</result>'
            )
        )

    def append_image_content(self, *, image_base64: str, image_info: str) -> None:
        """Append an image message to the context.

        Args:
            image_base64 (str): The base64-encoded image data.
            image_info (str): Additional information about the image (e.g., screen info).
        """
        self.append(
            ChatCompletionUserMessageParam(
                role='user',
                content=[
                    ChatCompletionContentPartImageParam(
                        type='image_url', image_url=ImageURL(url=f'data:image/png;base64,{image_base64}')
                    ),
                    ChatCompletionContentPartTextParam(type='text', text=f'** Screen Info **\n\n{image_info}'),
                ],
            )
        )

    def append_user_content(self, content: str) -> None:
        """Append a user message to the context.

        Args:
            content (str): The content of the user message.
        """
        self.append(ChatCompletionUserMessageParam(role='user', content=content))

    def reset(self) -> None:
        """Reset the context by clearing all messages."""
        self.clear()

    def remove_image_base64(self) -> None:
        """Remove base64 image data from the context to reduce size."""
        for message in self:
            if isinstance(message.get('content'), list):
                message['content'] = [item for item in message['content'] if item.get('type') == 'text']
