# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from __future__ import annotations

from base64 import b64decode, b64encode
from time import time
from uuid import uuid4

import orjson
from pydantic import BaseModel, Field

from .enums import MessageType, MimeType, ResponseStatus, Role, State
from .usage import Usage

ROOT_ID = "3713"


class Message(BaseModel):
    """
    Represents the content of the session both in request and response.
    """

    id: str = Field(
        default_factory=lambda: uuid4().hex[:12],
        title="Message ID",
        description="A unique identifier for the message - uuid4 int.",
    )

    parent_id: str = Field(
        ROOT_ID, title="Parent Message ID", description="The ID of the parent message."
    )

    turn: str = Field(
        default_factory=lambda: uuid4().hex[:12],
        title="Conversation Turn",
        description="The turn identifier is used to group/message are part of the same conversation turn.",
    )

    group: str = Field(
        default_factory=lambda: uuid4().hex,
        title="Message Group",
        description="The Group ID (UUID4) identify messages part of the same generation or action.",
    )

    actor: str = Field(
        "user", title="Actor", description="The actor of the message - user or agent."
    )

    role: Role = Field(
        Role.USER, title="Role", description="The role of the messages author."
    )

    # This is deprecated in favor of created_at, to allow for more precision. It will eventually be removed.
    timestamp: int = Field(
        default_factory=lambda: int(time()),
        title="Timestamp",
        description="DEPRECATED: The Unix timestamp (in seconds) of when the message was created.",
    )

    created_at: float = Field(
        default_factory=lambda: time(),
        title="Creation timestamp",
        description="The Unix timestamp (in seconds, including fractional parts) indicating when the message was created.",
    )

    message_type: MessageType = Field(
        ...,
        title="Message Type",
        description="The type of message - Generation, Tool Call, or Info.",
    )

    icon: str | None = Field(
        None, title="Icon", description="Custom svg icon to use in the UI."
    )

    text_color: str | None = Field(
        None, title="Text Color", description="Custom text color to use in the UI."
    )

    title: str | None = Field(
        None, title="Title", description="Descripting title of the message."
    )

    state: State = Field(
        State.START, title="State", description="The state the message belongs to."
    )

    content: str | None = Field(
        None,
        title="Message Content",
        description="The content of the message encoded as utf-8 bytes.",
    )

    mime_type: MimeType | None = Field(
        default=MimeType.TEXT,
        title="Content Type",
        description="The content type of the content field.",
    )

    status_code: int = Field(
        ResponseStatus.OK.value,
        title="Status Code",
        description="The status code of the message. 2xx is Okay, 4xx is a client error, 5xx is a server error.",
    )

    status_message: str = Field(
        ResponseStatus.OK.name,
        title="Status Message",
        description="Explain status code reason.",
    )

    usage: Usage | None = Field(
        default_factory=Usage,
        title="Model Usage Statistics",
        description="Token counts when message was generated by model.",
    )

    def to_json(self):
        "Returns the message as a dictionary."
        return orjson.dumps(self.model_dump())

    @staticmethod
    def from_json(data: dict | str) -> Message:
        "Creates a message from a json dictionary."
        if isinstance(data, str):
            data = orjson.loads(data)
        assert isinstance(data, dict)
        return Message(**data)

    def set_message_type(self, message_type: MessageType) -> Message:
        """
        Set the message type to a block message.
        """
        self.message_type = message_type
        self.set_content(message_type.value)
        return self

    def set_status(self, status: ResponseStatus) -> Message:
        """
        Set the status code and message of the message.
        """
        self.status_code = status.value
        self.status_message = status.name
        return self

    def set_content(
        self, content: str | bytes, mime_type: MimeType = MimeType.TEXT
    ) -> Message:
        """
        Set the content of the message.

        If `content` is of textual nature (as determined by the `mime_type`),
        the message's content is set to `content` itself, which is expected to
        be of type `str`. Otherwise, the message's content base64-encoded first.
        """
        if mime_type.value.startswith("text/"):
            assert isinstance(content, str)
            self.content = content
        else:
            assert isinstance(content, bytes)
            self.content = b64encode(content).decode("ascii")
        self.mime_type = mime_type
        return self

    def get_content(self) -> bytes | str:
        """
        Decodes the content of the message either to bytes or utf-8 string
        depending of the mime type.
        """
        if self.content is None:
            return ""

        assert self.mime_type is not None

        if self.mime_type.value.startswith("text/"):
            # Return text content as is
            return self.content
        else:
            # Base64 decode the content before returning it.
            return b64decode(self.content)
