#!/usr/bin/env python

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

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

from __future__ import annotations

import logging
import shutil
from enum import Enum
from pathlib import Path

from typing import TYPE_CHECKING

from vi import Client
from vi.api.resources.datasets.assets.responses import FailureMode
from vi.api.resources.datasets.assets.results import AssetUploadResult
from vi.utils.graceful_exit import GracefulExit

if TYPE_CHECKING:
    from vi.utils.progress import ViProgress

logger = logging.getLogger(__name__)


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

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


def upload_batch(
    client: Client,
    batch_file_paths: list[Path],
    batch_temp_dir: Path,
    batch_num: int,
    handler: GracefulExit,
    dataset_id: str,
    failure_mode: FailureMode,
    on_asset_overwritten: str,
    wait_until_done: bool,
    progress: ViProgress | None = None,
) -> AssetUploadResult | None:
    """Upload a batch of images from disk.

    This function receives file paths to images saved on disk and uploads them.
    After upload, the temporary files are cleaned up to free disk space.

    Progress is displayed by the asset uploader's native progress bar.

    Args:
        client: Client instance
        batch_file_paths: List of Path objects pointing to saved images
        batch_temp_dir: Temporary directory containing the batch images
        batch_num: Batch number (for logging)
        handler: Graceful exit handler
        dataset_id: Target dataset ID
        failure_mode: Failure mode enum
        on_asset_overwritten: Action when asset exists ("RemoveLinkedResource" or "KeepLinkedResource")
        wait_until_done: Whether to wait for upload processing to complete
        progress: Optional ViProgress instance to use for progress display

    Returns:
        AssetUploadResult for this batch, or None if cancelled

    """
    if handler.exit_now:
        return None

    logger.debug(f"Processing batch {batch_num}: {len(batch_file_paths)} assets")

    if handler.exit_now:
        logger.warning(f"Batch {batch_num}: Upload cancelled by user")
        # Clean up temp directory even if cancelled
        try:
            if batch_temp_dir.exists():
                shutil.rmtree(batch_temp_dir)
                logger.debug(f"Cleaned up temp directory: {batch_temp_dir}")
        except Exception as e:
            logger.warning(f"Failed to clean up temp directory {batch_temp_dir}: {e}")
        return None

    if not batch_file_paths:
        logger.error(f"Batch {batch_num}: No assets provided!")
        return None

    try:
        # Upload batch using file paths
        # Use asset uploader's native progress display
        logger.info(f"Uploading batch {batch_num}: {len(batch_file_paths)} assets")
        result = client.assets.upload(
            dataset_id=dataset_id,
            paths=batch_file_paths,
            failure_mode=failure_mode.value,
            on_asset_overwritten=on_asset_overwritten,
            wait_until_done=wait_until_done,
            show_progress=progress is not None,
            progress=progress,
        )

        logger.info(
            f"Batch {batch_num} complete: {result.total_succeeded} succeeded, "
            f"{result.total_failed} failed"
        )

        return result

    finally:
        # Clean up temporary directory after upload
        try:
            if batch_temp_dir.exists():
                shutil.rmtree(batch_temp_dir)
                logger.debug(f"Cleaned up temp directory: {batch_temp_dir}")
        except Exception as e:
            logger.warning(f"Failed to clean up temp directory {batch_temp_dir}: {e}")
