#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   messages.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK inference messages module.
"""

from vi.inference.utils.media_utils import decode_base64_image, decode_base64_video


def create_system_message(text: str) -> dict:
    """Create a system message with text content.

    Args:
        text: The system prompt text.

    Returns:
        Dictionary representing a system message with string content.

    """
    return {
        "role": "system",
        "content": [
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_image(image_path: str, text: str) -> dict:
    """Create a user message with image and text content.

    Args:
        image_path: Path to the image file. For URLs or data URIs,
            use create_user_message_with_image_url() instead.
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with image and text content.

    """
    return {
        "role": "user",
        "content": [
            {"type": "image", "path": image_path},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_image_url(image_url: str, text: str) -> dict:
    """Create a user message with image URL and text content.

    Supports HTTP(S) URLs and data URIs. For data URIs, decodes to PIL Image
    and includes the object directly in the message.

    Args:
        image_url: URL or data URI of the image
            (e.g., "http://example.com/image.jpg" or "data:image/jpeg;base64,...").
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with image content and text.

    """
    # If it's a data URI, decode to PIL Image object
    if image_url.startswith("data:image"):
        image_object = decode_base64_image(image_url)
        return {
            "role": "user",
            "content": [
                {"type": "image", "image": image_object},
                {"type": "text", "text": text},
            ],
        }

    # For regular URLs, include the URL
    return {
        "role": "user",
        "content": [
            {"type": "image", "url": image_url},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_base64_image(base64_image: str, text: str) -> dict:
    """Create a user message with base64-encoded image and text content.

    Converts the base64 string to a PIL Image object for proper processing.

    Args:
        base64_image: Base64-encoded image string. Can be raw base64
            (e.g., "/9j/4AAQSkZJRgABAQEAYABgAAD...") or a data URI
            (e.g., "data:image/jpeg;base64,/9j/...").
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with image object and text content.

    Raises:
        ValueError: If the base64 string cannot be decoded.

    """
    # Decode base64 to PIL Image using helper function
    image_object = decode_base64_image(base64_image)

    return {
        "role": "user",
        "content": [
            {"type": "image", "image": image_object},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_video(video_path: str, text: str) -> dict:
    """Create a user message with video and text content.

    Args:
        video_path: Path to the video file, or list of frame image paths.
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with video and text content.

    Note:
        The fps parameter for video sampling should be passed to
        processor.apply_chat_template() instead of in the message content.
        For example: processor.apply_chat_template(messages, fps=4.0)

    """
    return {
        "role": "user",
        "content": [
            {"type": "video", "path": video_path},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_video_url(video_url: str, text: str) -> dict:
    """Create a user message with video URL and text content.

    Supports HTTP(S) URLs and data URIs. For data URIs, decodes to numpy array
    and includes the object directly in the message.

    Args:
        video_url: URL or data URI of the video
            (e.g., "http://example.com/video.mp4" or "data:video/mp4;base64,...").
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with video content and text.

    """
    # If it's a data URI, decode to numpy array
    if video_url.startswith("data:video"):
        video_object = decode_base64_video(video_url)
        return {
            "role": "user",
            "content": [
                {"type": "video", "video": video_object},
                {"type": "text", "text": text},
            ],
        }

    # For regular URLs, include the URL
    return {
        "role": "user",
        "content": [
            {"type": "video", "url": video_url},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_base64_video(base64_video: str, text: str) -> dict:
    """Create a user message with base64-encoded video and text content.

    Converts the base64 string to video frames for proper processing.

    Args:
        base64_video: Base64-encoded video string. Can be raw base64 or a data URI
            (e.g., "data:video/mp4;base64,...").
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with video object and text content.

    Raises:
        ImportError: If transformers video utilities are not available.
        ValueError: If the base64 string cannot be decoded.

    """
    # Decode base64 to video frames using helper function
    video_object = decode_base64_video(base64_video)

    return {
        "role": "user",
        "content": [
            {"type": "video", "video": video_object},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_image_object(image_object: object, text: str) -> dict:
    """Create a user message with an image object and text content.

    Pass PIL Image or numpy array directly (useful for preprocessed images).

    Args:
        image_object: PIL Image object or numpy array (shape: [height, width, channels]).
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with image object and text content.

    """
    return {
        "role": "user",
        "content": [
            {"type": "image", "image": image_object},
            {"type": "text", "text": text},
        ],
    }


def create_user_message_with_video_object(video_object: object, text: str) -> dict:
    """Create a user message with a video object and text content.

    Pass already-decoded video frames as numpy array (useful for preprocessed videos).

    Args:
        video_object: Numpy array with shape (num_frames, height, width, channels).
            Useful when you've preprocessed frames with OpenCV, decord, or torchvision.
        text: The user prompt text.

    Returns:
        Dictionary representing a user message with video object and text content.

    """
    return {
        "role": "user",
        "content": [
            {"type": "video", "video": video_object},
            {"type": "text", "text": text},
        ],
    }
