#!/usr/bin/env python

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

@File    :   batch_processor.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   HuggingFace dataset batch processor for parallel upload processing.
"""

from __future__ import annotations

import logging
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from threading import Lock
from typing import Any

from PIL import Image
from vi import Client
from vi.api.resources.datasets.assets.results import AssetUploadResult
from vi.dataset.loaders.hf import consts
from vi.dataset.loaders.hf.uploader import ImageFormat, upload_batch
from vi.dataset.loaders.hf.utils import extract_image_from_example, save_image_to_disk
from vi.utils.graceful_exit import GracefulExit
from vi.utils.progress import ViProgress

logger = logging.getLogger(__name__)


def generate_batches(
    dataset: Any,
    image_column: str,
    batch_size: int,
    limit: int | None,
    handler: GracefulExit,
    image_format: ImageFormat,
    progress: ViProgress | None,
    total_images: int,
):
    """Generate batches of image file paths by streaming and saving images to disk in parallel.

    This implementation saves images to a temporary directory before upload:
    1. Extracts up to 16 images at a time
    2. Saves them to disk in parallel with 16 worker threads
    3. Closes images to free memory after saving
    4. Yields file paths for batch upload

    Args:
        dataset: Dataset to stream from
        image_column: Name of the column containing image data
        batch_size: Maximum images per batch
        limit: Maximum total images to process
        handler: Graceful exit handler for cancellation
        image_format: Image format for saving (png, jpeg, webp)
        progress: Progress tracker for UI updates
        total_images: Total number of images (for progress tracking)

    Yields:
        Tuples of (batch_num, batch_temp_dir, batch_file_paths) where:
        - batch_num: The batch number (0-indexed)
        - batch_temp_dir: Temporary directory containing the batch images
        - batch_file_paths: List of Path objects pointing to saved images

    """
    batch_file_paths: list[Path] = []
    batch_num = 0
    total_processed = 0
    seen_filenames: set[str] = set()
    seen_filenames_lock = Lock()

    # Create a temporary directory for this entire upload session
    session_temp_dir = Path(tempfile.mkdtemp(prefix="hf_dataset_upload_"))
    logger.info(f"Created temporary directory: {session_temp_dir}")

    # Create progress task for streaming and saving phase
    conversion_progress_task = None
    if progress:
        conversion_progress_task = progress.add_task(
            "Streaming and saving images to disk...",
            total=total_images,
        )

    try:
        # Iterator for dataset
        dataset_iter = iter(dataset)

        while True:
            # Check for cancellation
            if handler.exit_now:
                logger.info("Batch generation cancelled by user")
                if progress and conversion_progress_task is not None:
                    progress.update(
                        conversion_progress_task, description="✗ Streaming cancelled"
                    )
                break

            # Check global limit
            if limit is not None and total_processed >= limit:
                logger.debug(f"Reached limit of {limit} images")
                break

            # Create batch-specific temporary directory
            batch_temp_dir = session_temp_dir / f"batch_{batch_num}"
            batch_temp_dir.mkdir(parents=True, exist_ok=True)

            # Extract up to MAX_PARALLEL_PROCESSING_WORKERS images
            images_to_save: list[tuple[Image.Image, dict[str, Any], int]] = []

            for _ in range(consts.MAX_PARALLEL_PROCESSING_WORKERS):
                try:
                    example = next(dataset_iter)

                    # Check limit again
                    if (
                        limit is not None
                        and (total_processed + len(images_to_save)) >= limit
                    ):
                        break

                    # Extract image
                    image = extract_image_from_example(example, image_column)

                    if image is not None:
                        images_to_save.append(
                            (image, example, total_processed + len(images_to_save))
                        )

                except StopIteration:
                    break

            # If no images extracted, we're done
            if not images_to_save:
                break

            # Save images to disk in parallel using thread pool
            with ThreadPoolExecutor(
                max_workers=consts.MAX_PARALLEL_PROCESSING_WORKERS
            ) as executor:

                def save_single_image(
                    img_data: tuple[Image.Image, dict[str, Any], int],
                ) -> Path | None:
                    image, metadata, image_index = img_data
                    return save_image_to_disk(
                        image=image,
                        metadata=metadata,
                        image_column=image_column,
                        image_index=image_index,
                        image_format=image_format,
                        temp_dir=batch_temp_dir,
                        seen_filenames=seen_filenames,
                        seen_filenames_lock=seen_filenames_lock,
                    )

                # Submit all save tasks
                futures = [
                    executor.submit(save_single_image, img_data)
                    for img_data in images_to_save
                ]

                # Collect results as they complete
                for future in as_completed(futures):
                    if handler.exit_now:
                        break

                    try:
                        file_path = future.result()

                        if file_path:
                            batch_file_paths.append(file_path)
                            total_processed += 1

                            # Update progress after each image is saved
                            if progress and conversion_progress_task is not None:
                                progress.update(
                                    conversion_progress_task,
                                    description=f"Saved {total_processed} / {total_images} images to disk",
                                    advance=1,
                                )

                            # Yield batch when it reaches batch_size
                            if len(batch_file_paths) >= batch_size:
                                yield (batch_num, batch_temp_dir, batch_file_paths)
                                batch_num += 1
                                batch_file_paths = []  # Clear list for next batch
                                with seen_filenames_lock:
                                    seen_filenames.clear()  # Clear for next batch
                                # Create new batch directory for next batch
                                batch_temp_dir = session_temp_dir / f"batch_{batch_num}"
                                batch_temp_dir.mkdir(parents=True, exist_ok=True)

                    except Exception as e:
                        logger.warning(f"Failed to save image in parallel: {e}")

        # Yield remaining assets in final batch
        if batch_file_paths and not handler.exit_now:
            yield (batch_num, batch_temp_dir, batch_file_paths)

        # Mark conversion phase as complete
        if progress and conversion_progress_task is not None and not handler.exit_now:
            progress.update(
                conversion_progress_task,
                description=f"✓ Saved {total_processed} images to disk",
            )

    finally:
        # Remove conversion progress task after a brief delay so user can see completion
        if progress and conversion_progress_task is not None:
            time.sleep(0.5)  # Brief delay to show completion message
            progress.remove_task(conversion_progress_task)


class HFBatchProcessor:
    """Handles parallel processing of HF dataset batches.

    This class manages the execution of multiple upload batches in parallel while
    respecting the configured parallelism limits. It implements dynamic batch
    submission to ensure only a limited number of batches are active at once.

    Attributes:
        client: Client for uploading assets
        config: Configuration dict containing upload parameters
        handler: Graceful exit handler for cancellation support
        logger: Logger for this processor

    """

    def __init__(
        self,
        client: Client,
        config: dict[str, Any],
        progress: ViProgress | None,
        handler: GracefulExit,
    ):
        """Initialize the parallel batch processor.

        Args:
            client: Client for uploading assets
            config: Configuration dict with upload parameters
            progress: Progress tracker for UI updates
            handler: Handler for graceful cancellation

        """
        self.client = client
        self.config = config
        self.progress = progress
        self.handler = handler
        self.logger = logging.getLogger("hf.loader.parallel")

    def process_batches_parallel(
        self,
        batch_generator: Any,
        max_parallel: int,
    ) -> list[AssetUploadResult]:
        """Process batches in parallel with controlled concurrency.

        Upload progress is displayed by the asset uploader for each batch.

        Args:
            batch_generator: Generator yielding (batch_num, batch_images) tuples
            max_parallel: Maximum number of batches to process in parallel

        Returns:
            List of successful upload results

        """
        session_results = []

        executor = ThreadPoolExecutor(max_workers=max_parallel)
        try:
            future_to_batch = {}

            # Submit initial batches up to max_parallel
            for _ in range(max_parallel):
                if not self._submit_next_batch(
                    executor, batch_generator, future_to_batch
                ):
                    break  # No more batches available

            # Process completions and submit new batches to maintain parallelism
            while future_to_batch:
                if self.handler.exit_now:
                    self._cancel_all_batches(future_to_batch)
                    break

                try:
                    result = self._wait_for_next_completion(
                        future_to_batch, batch_generator, executor
                    )

                    if result:
                        session_results.append(result)
                except TimeoutError:
                    # No future completed within timeout, loop again to check exit_now
                    continue

        finally:
            # Ensure executor is shut down immediately and don't wait for workers
            executor.shutdown(wait=False, cancel_futures=True)

        return session_results

    def _submit_initial_batches(
        self,
        executor: ThreadPoolExecutor,
        batch_generator: Any,
        future_to_batch: dict,
        count: int,
    ):
        """Submit the initial wave of batches for processing.

        Args:
            executor: Thread pool executor for parallel processing
            batch_generator: Generator yielding batch data
            future_to_batch: Dictionary mapping futures to batch numbers
            count: Number of initial batches to submit

        """
        for _ in range(count):
            try:
                batch_num, batch_temp_dir, batch_file_paths = next(batch_generator)
                future = self._submit_batch(
                    executor,
                    batch_num,
                    batch_file_paths,
                    batch_temp_dir,
                )
                future_to_batch[future] = batch_num
            except StopIteration:
                break

    def _submit_batch(
        self,
        executor: ThreadPoolExecutor,
        batch_num: int,
        batch_file_paths: list[Path],
        batch_temp_dir: Path,
    ):
        """Submit a single batch for processing.

        Args:
            executor: Thread pool executor
            batch_num: Index of this batch (0-based)
            batch_file_paths: List of Path objects pointing to saved images
            batch_temp_dir: Temporary directory containing the batch images

        Returns:
            Future representing the batch processing task

        """
        return executor.submit(
            upload_batch,
            client=self.client,
            batch_file_paths=batch_file_paths,
            batch_temp_dir=batch_temp_dir,
            batch_num=batch_num,
            handler=self.handler,
            progress=self.progress,
            **self.config,
        )

    def _wait_for_next_completion(
        self,
        future_to_batch: dict,
        batch_generator: Any,
        executor: ThreadPoolExecutor,
    ) -> AssetUploadResult | None:
        """Wait for the next batch to complete and submit a new one if available.

        Args:
            future_to_batch: Dictionary mapping futures to batch numbers
            batch_generator: Generator over remaining batches
            executor: Thread pool executor

        Returns:
            The upload result from the completed batch, or None if no batches completed

        """
        # Use timeout in as_completed to allow checking for cancellation
        for future in as_completed(future_to_batch, timeout=0.1):
            # Check for cancellation before processing result
            if self.handler.exit_now:
                return None

            batch_num = future_to_batch[future]

            try:
                result = future.result()
                del future_to_batch[future]

                # Submit next batch to maintain parallelism level
                self._submit_next_batch(executor, batch_generator, future_to_batch)

                return result

            except Exception as e:
                self.logger.error(
                    f"Batch {batch_num} failed: {e}. Cancelling remaining batches."
                )
                self._cancel_all_batches(future_to_batch)
                future_to_batch.clear()
                raise

        return None

    def _submit_next_batch(
        self,
        executor: ThreadPoolExecutor,
        batch_generator: Any,
        future_to_batch: dict,
    ) -> bool:
        """Submit the next batch if one is available.

        Args:
            executor: Thread pool executor
            batch_generator: Generator over remaining batches
            future_to_batch: Dictionary mapping futures to batch numbers

        Returns:
            True if a batch was submitted, False if no more batches available

        """
        try:
            batch_num, batch_temp_dir, batch_file_paths = next(batch_generator)
            future = self._submit_batch(
                executor, batch_num, batch_file_paths, batch_temp_dir
            )
            future_to_batch[future] = batch_num
            return True
        except StopIteration:
            return False

    def _cancel_all_batches(self, future_to_batch: dict):
        """Cancel all pending batches.

        Args:
            future_to_batch: Dictionary of futures to cancel

        """
        for future in future_to_batch:
            future.cancel()

    def process_batches_sequential(
        self,
        batch_generator: Any,
    ) -> list[AssetUploadResult]:
        """Process batches sequentially one at a time.

        Upload progress is displayed by the asset uploader for each batch.

        Args:
            batch_generator: Generator yielding batches

        Returns:
            List of upload results

        """
        results = []

        try:
            for batch_num, batch_temp_dir, batch_file_paths in batch_generator:
                if self.handler.exit_now:
                    break

                result = upload_batch(
                    client=self.client,
                    batch_file_paths=batch_file_paths,
                    batch_temp_dir=batch_temp_dir,
                    batch_num=batch_num,
                    handler=self.handler,
                    progress=self.progress,
                    **self.config,
                )

                if result:
                    results.append(result)

        finally:
            pass

        return results
