mono/artdag/common/artdag_common/utils/media.py

"""
Media type detection and handling utilities.
"""

from pathlib import Path
from typing import Optional
import mimetypes

# Initialize mimetypes database
mimetypes.init()

# Media type categories
VIDEO_TYPES = {"video/mp4", "video/webm", "video/quicktime", "video/x-msvideo", "video/avi"}
IMAGE_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp", "image/svg+xml"}
AUDIO_TYPES = {"audio/mpeg", "audio/wav", "audio/ogg", "audio/flac", "audio/aac", "audio/mp3"}

# File extension mappings
EXTENSION_TO_CATEGORY = {
    # Video
    ".mp4": "video",
    ".webm": "video",
    ".mov": "video",
    ".avi": "video",
    ".mkv": "video",
    # Image
    ".jpg": "image",
    ".jpeg": "image",
    ".png": "image",
    ".gif": "image",
    ".webp": "image",
    ".svg": "image",
    # Audio
    ".mp3": "audio",
    ".wav": "audio",
    ".ogg": "audio",
    ".flac": "audio",
    ".aac": "audio",
    ".m4a": "audio",
}


def detect_media_type(path: Path) -> str:
    """
    Detect the media category for a file.

    Args:
        path: Path to the file

    Returns:
        Category string: "video", "image", "audio", or "unknown"
    """
    if not path:
        return "unknown"

    # Try extension first
    ext = path.suffix.lower()
    if ext in EXTENSION_TO_CATEGORY:
        return EXTENSION_TO_CATEGORY[ext]

    # Try mimetypes
    mime_type, _ = mimetypes.guess_type(str(path))
    if mime_type:
        if mime_type in VIDEO_TYPES or mime_type.startswith("video/"):
            return "video"
        if mime_type in IMAGE_TYPES or mime_type.startswith("image/"):
            return "image"
        if mime_type in AUDIO_TYPES or mime_type.startswith("audio/"):
            return "audio"

    return "unknown"


def get_mime_type(path: Path) -> str:
    """
    Get the MIME type for a file.

    Args:
        path: Path to the file

    Returns:
        MIME type string or "application/octet-stream"
    """
    mime_type, _ = mimetypes.guess_type(str(path))
    return mime_type or "application/octet-stream"


def get_media_extension(media_type: str) -> str:
    """
    Get the typical file extension for a media type.

    Args:
        media_type: Media category or MIME type

    Returns:
        File extension with dot (e.g., ".mp4")
    """
    if media_type == "video":
        return ".mp4"
    if media_type == "image":
        return ".png"
    if media_type == "audio":
        return ".mp3"

    # Try as MIME type
    ext = mimetypes.guess_extension(media_type)
    return ext or ""


def is_streamable(path: Path) -> bool:
    """
    Check if a file type is streamable (video/audio).

    Args:
        path: Path to the file

    Returns:
        True if the file can be streamed
    """
    media_type = detect_media_type(path)
    return media_type in ("video", "audio")


def needs_conversion(path: Path, target_format: str = "mp4") -> bool:
    """
    Check if a video file needs format conversion.

    Args:
        path: Path to the file
        target_format: Target format (default mp4)

    Returns:
        True if conversion is needed
    """
    media_type = detect_media_type(path)
    if media_type != "video":
        return False

    ext = path.suffix.lower().lstrip(".")
    return ext != target_format


def get_video_src(
    content_hash: str,
    original_path: Optional[Path] = None,
    is_ios: bool = False,
) -> str:
    """
    Get the appropriate video source URL.

    For iOS devices, prefer MP4 format.

    Args:
        content_hash: Content hash for the video
        original_path: Optional original file path
        is_ios: Whether the client is iOS

    Returns:
        URL path for the video source
    """
    if is_ios:
        return f"/cache/{content_hash}/mp4"

    if original_path and original_path.suffix.lower() in (".mp4", ".webm"):
        return f"/cache/{content_hash}/raw"

    return f"/cache/{content_hash}/mp4"