from typing import Any, Literal

from PIL.Image import Image
from pydantic import BaseModel, Field

from screensuite.resize_image import resize_image


class ImageContent(BaseModel):
    """Represents an image content."""

    type: Literal["image"] = "image"
    image: Image

    class Config:
        arbitrary_types_allowed = True

    def model_post_init(self, context: Any) -> None:
        self.image = resize_image(self.image)[0]


class TextContent(BaseModel):
    """Represents a text content."""

    type: Literal["text"] = "text"
    text: str


MessageContent = ImageContent | TextContent


class ChatMessage(BaseModel):
    """Represents a message to the model."""

    role: Literal["system", "user", "assistant"]
    content: list[MessageContent] = Field(default_factory=list)


def dump_chat_message_list(content: list[ChatMessage]) -> list[dict]:
    """Get a list of dicts from a list of ChatMessage."""
    return [message.model_dump() for message in content]
