#!/usr/bin/env python

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

@File    :   utils.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   HuggingFace dataset utils module.
"""

import io
import logging
import threading
from enum import Enum
from pathlib import Path
from typing import Any

from PIL import Image

logger = logging.getLogger(__name__)


class ImageFormat(Enum):
    """Image format enum for saving images."""

    PNG = "png"
    JPEG = "jpeg"
    WEBP = "webp"


def extract_image_from_example(
    example: dict[str, Any], image_column: str
) -> Image.Image | None:
    """Extract PIL Image from HuggingFace dataset example.

    Args:
        example: Dataset example dictionary.
        image_column: Name of the column containing image data.

    Returns:
        PIL Image object, or None if extraction fails.

    Raises:
        ValueError: If image column doesn't exist.

    """
    if image_column not in example:
        available_columns = list(example.keys())
        msg = (
            f"Image column '{image_column}' not found in dataset. "
            f"Available columns: {available_columns}"
        )
        raise ValueError(msg)

    image_data = example[image_column]

    # Handle different image formats
    try:
        image = None

        # Case 1: Already a PIL Image
        if isinstance(image_data, Image.Image):
            image = image_data

        # Case 2: Bytes or BytesIO
        elif isinstance(image_data, (bytes, bytearray)):
            image = Image.open(io.BytesIO(image_data))

        # Case 3: File path (string or Path)
        elif isinstance(image_data, (str, Path)):
            image = Image.open(image_data)

        # Case 4: Dict with 'bytes' or 'path' key (HF Datasets format)
        elif isinstance(image_data, dict):
            if "bytes" in image_data:
                image = Image.open(io.BytesIO(image_data["bytes"]))
            elif "path" in image_data:
                image = Image.open(image_data["path"])

        if image is None:
            # Unable to extract image
            logger.warning(
                f"Unable to extract image from column '{image_column}'. "
                f"Type: {type(image_data)}"
            )
            return None

        return image

    except Exception as e:
        logger.warning(f"Failed to load image: {e}")
        return None


def extract_filename_from_metadata(
    metadata: dict[str, Any], image_column: str
) -> str | None:
    """Extract original filename from dataset metadata.

    Args:
        metadata: Dataset example metadata dictionary
        image_column: Name of the image column

    Returns:
        Original filename if found, None otherwise

    """
    # Common metadata keys that might contain the filename
    filename_keys = [
        "file_name",
        "filename",
        "image_name",
        "image_file",
        "path",
        "image_path",
        "file",
    ]

    for key in filename_keys:
        if key in metadata:
            value = metadata[key]
            if isinstance(value, str) and value:
                # Extract just the filename if it's a full path
                return Path(value).name
            if isinstance(value, Path):
                return value.name

    # Check if the image column itself has path information
    if image_column in metadata and isinstance(metadata[image_column], dict):
        img_dict = metadata[image_column]
        if "path" in img_dict and img_dict["path"]:
            return Path(img_dict["path"]).name

    return None


def save_image_to_disk(
    image: Image.Image,
    metadata: dict[str, Any],
    image_column: str,
    image_index: int,
    image_format: ImageFormat,
    temp_dir: Path,
    seen_filenames: set[str],
    seen_filenames_lock: type[threading.Lock] | None = None,
) -> Path | None:
    """Save a PIL Image to disk for upload.

    Args:
        image: PIL Image to save
        metadata: Dataset metadata for filename extraction
        image_column: Column name containing image data
        image_index: Global image index for fallback filenames
        image_format: Format to save image as (png, jpeg, webp)
        temp_dir: Temporary directory to save images to
        seen_filenames: Set of already used filenames in this batch
        seen_filenames_lock: Optional lock for thread-safe access to seen_filenames

    Returns:
        Path to saved image file, or None if save failed

    """
    try:
        # Try to extract original filename from metadata
        original_filename = extract_filename_from_metadata(metadata, image_column)

        if original_filename:
            # Use original filename, ensure unique (thread-safe)
            base_name = Path(original_filename).stem
            extension = Path(original_filename).suffix or f".{image_format.value}"
            filename = original_filename
            counter = 1

            if seen_filenames_lock:
                with seen_filenames_lock:
                    while filename in seen_filenames:
                        filename = f"{base_name}_{counter}{extension}"
                        counter += 1
                    seen_filenames.add(filename)
            else:
                while filename in seen_filenames:
                    filename = f"{base_name}_{counter}{extension}"
                    counter += 1
                seen_filenames.add(filename)
        else:
            # Fallback to indexed filename
            filename = f"image_{image_index:06d}.{image_format.value}"

        # Save image to disk
        file_path = temp_dir / filename
        image.save(file_path, format=image_format.value.upper())

        return file_path

    except Exception as e:
        logger.warning(f"Failed to save image {image_index} to disk: {e}")
        return None
