blob_utils.py

"""
Blob conversion utilities for Gradio image components.
Handles conversion of blob data to proper image file formats.
"""

import hashlib
import os
from typing import Dict, Any


class BlobConverter:
    """Handles conversion of blob data to proper image file formats."""

    # File format signatures
    FORMAT_SIGNATURES = {
        b"\x89PNG\r\n\x1a\n": (".png", "image/png"),
        b"\xff\xd8\xff": (".jpg", "image/jpeg"),
        b"GIF87a": (".gif", "image/gif"),
        b"GIF89a": (".gif", "image/gif"),
    }

    @classmethod
    def is_blob_data(cls, image_data: Dict[str, Any]) -> bool:
        """
        Check if the image data represents a blob that needs conversion.

        Args:
            image_data: Dictionary containing image metadata

        Returns:
            True if the data is a blob that needs conversion
        """
        return (
            isinstance(image_data, dict)
            and "path" in image_data
            and "blob" in image_data["path"]
            and image_data.get("size") is None
            and image_data.get("orig_name") is None
            and image_data.get("mime_type") is None
        )

    @classmethod
    def detect_format(cls, content: bytes) -> tuple[str, str]:
        """
        Detect image format from file content.

        Args:
            content: Binary content of the file

        Returns:
            Tuple of (extension, mime_type)
        """
        for signature, (ext, mime_type) in cls.FORMAT_SIGNATURES.items():
            if content.startswith(signature):
                return ext, mime_type

        # Default to PNG if format cannot be determined
        return ".png", "image/png"

    @classmethod
    def generate_filename(cls, content: bytes, extension: str) -> str:
        """
        Generate a unique filename for the converted blob.

        Args:
            content: Binary content of the file
            extension: File extension to use

        Returns:
            Unique filename
        """
        content_hash = hashlib.md5(content).hexdigest()[:8]
        return f"flagged_image_{content_hash}{extension}"

    @classmethod
    def convert_blob(cls, image_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert blob data to proper image file format.

        Args:
            image_data: Original image data dictionary

        Returns:
            Updated image data with proper file information
        """
        if not cls.is_blob_data(image_data):
            return image_data

        print(f"DEBUG: Converting blob data: {image_data}")

        blob_path = image_data["path"]
        print(f"DEBUG: Blob path: {blob_path}")

        # Read blob content
        with open(blob_path, "rb") as f:
            content = f.read()

        file_size = len(content)
        print(f"DEBUG: File size: {file_size}")

        # Detect format
        extension, mime_type = cls.detect_format(content)
        print(f"DEBUG: Detected format: {extension}, MIME type: {mime_type}")

        # Generate filename and path
        filename = cls.generate_filename(content, extension)
        temp_dir = os.path.dirname(blob_path)
        new_path = os.path.join(temp_dir, filename)

        print(f"DEBUG: Generated filename: {filename}")
        print(f"DEBUG: New path: {new_path}")

        # Write converted file
        with open(new_path, "wb") as f:
            f.write(content)

        print(f"DEBUG: Successfully converted blob to: {new_path}")

        # Return updated image data
        converted_data = {
            "path": new_path,
            "url": image_data["url"].replace("blob", filename),
            "size": file_size,
            "orig_name": filename,
            "mime_type": mime_type,
            "is_stream": False,
            "meta": image_data.get("meta", {}),
        }

        print(f"DEBUG: Converted data: {converted_data}")
        return converted_data


def decode_blob_data(image_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Convenience function to decode blob data from Gradio image component.

    Args:
        image_data: Image data dictionary from Gradio

    Returns:
        Converted image data or original data if not a blob
    """
    print(f"DEBUG: Original input - image: {image_data}")

    result = BlobConverter.convert_blob(image_data)

    if result is image_data:
        print("DEBUG: Not a blob, skipping conversion")
    else:
        print("DEBUG: Blob conversion completed")

    return result


def is_blob_data(image_data: Dict[str, Any]) -> bool:
    """
    Check if the image data represents a blob that needs conversion.

    Args:
        image_data: Dictionary containing image metadata

    Returns:
        True if the data is a blob that needs conversion
    """
    return BlobConverter.is_blob_data(image_data)