#!/usr/bin/env python

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

@File    :   processor.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Parallel batch processor for asset uploads.
"""

from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import TYPE_CHECKING

from vi.api.resources.datasets.assets import responses
from vi.api.resources.datasets.assets.types import AssetSource, BatchConfig
from vi.logging import get_logger

if TYPE_CHECKING:
    from vi.api.resources.datasets.assets.uploader import (
        AssetUploader,
        AssetUploadProgressTracker,
    )
    from vi.utils.graceful_exit import GracefulExit
    from vi.utils.progress import ViProgress

logger = get_logger("assets.processor")


class ParallelBatchProcessor:
    """Handles parallel processing of upload 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:
        uploader: The AssetUploader instance that owns this processor
        config: Batch configuration containing upload parameters
        progress: Progress display for tracking upload status
        handler: Graceful exit handler for cancellation support
        master_tracker: Master progress tracker shared across all batches

    """

    def __init__(
        self,
        uploader: "AssetUploader",
        config: BatchConfig,
        progress: "ViProgress | None",
        handler: "GracefulExit",
        master_tracker: "AssetUploadProgressTracker | None" = None,
    ):
        """Initialize the parallel batch processor.

        Args:
            uploader: The AssetUploader instance
            config: Configuration for batch uploads
            progress: Progress tracker for UI updates
            handler: Handler for graceful cancellation
            master_tracker: Optional master progress tracker for all batches

        """
        self.uploader = uploader
        self.config = config
        self.progress = progress
        self.handler = handler
        self.master_tracker = master_tracker

    def process_batches(
        self,
        batches: list[list[AssetSource]],
        max_parallel: int,
    ) -> list[responses.AssetIngestionSession]:
        """Process batches with controlled parallelism.

        Args:
            batches: List of file path batches
            max_parallel: Maximum number of batches to process in parallel

        Returns:
            List of successful ingestion session responses

        """
        total_batches = len(batches)
        session_responses = []

        executor = ThreadPoolExecutor(max_workers=max_parallel)
        try:
            future_to_batch = {}
            batch_iter = enumerate(batches, start=1)

            # Submit initial batches
            self._submit_initial_batches(
                executor,
                batch_iter,
                future_to_batch,
                min(max_parallel, total_batches),
                total_batches,
            )

            # Process completions and submit new batches
            # When cancellation is detected, we still wait for all batch workers to finish
            # so they can properly abort ingestion sessions and clean up resources
            cancellation_detected = False

            while future_to_batch:
                if self.handler.exit_now and not cancellation_detected:
                    cancellation_detected = True
                    logger.info(
                        f"Cancellation detected, waiting for {len(future_to_batch)} "
                        f"batch workers to finish cleanup"
                    )
                    # Cancel futures that haven't started yet
                    self._cancel_all_batches(future_to_batch)

                try:
                    session_response = self._wait_for_next_completion(
                        future_to_batch, batch_iter, executor, total_batches
                    )

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

        finally:
            # Wait for all batch workers to finish cleanup before shutting down
            # This ensures ingestion sessions are properly aborted
            executor.shutdown(wait=True, cancel_futures=True)

        return session_responses

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

        This submits up to `count` batches to the thread pool executor to begin
        parallel processing. This is the first step in the dynamic submission strategy.

        Args:
            executor: Thread pool executor for parallel processing
            batch_iter: Iterator over (batch_idx, batch_assets) tuples
            future_to_batch: Dictionary mapping futures to batch indices
            count: Number of initial batches to submit
            total_batches: Total number of batches (for progress tracking)

        """
        for _ in range(count):
            try:
                batch_idx, batch_assets = next(batch_iter)
                future = self._submit_batch(
                    executor, batch_idx, batch_assets, total_batches
                )
                future_to_batch[future] = batch_idx
            except StopIteration:
                break

    def _submit_batch(
        self,
        executor: ThreadPoolExecutor,
        batch_idx: int,
        batch_assets: list[AssetSource],
        total_batches: int,
    ):
        """Submit a single batch for processing.

        Args:
            executor: Thread pool executor
            batch_idx: Index of this batch (1-based)
            batch_assets: List of asset sources in this batch
            total_batches: Total number of batches

        Returns:
            Future representing the batch processing task

        """
        return executor.submit(
            self.uploader.process_single_batch,
            config=self.config,
            batch_assets=batch_assets,
            batch_idx=batch_idx,
            total_batches=total_batches,
            progress=self.progress,
            handler=self.handler,
            master_tracker=self.master_tracker,
        )

    def _wait_for_next_completion(
        self,
        future_to_batch: dict,
        batch_iter: enumerate,
        executor: ThreadPoolExecutor,
        total_batches: int,
    ) -> responses.AssetIngestionSession | None:
        """Wait for the next batch to complete and submit a new one if available.

        This implements the dynamic submission strategy: as soon as one batch completes,
        a new batch is submitted to keep the thread pool fully utilized.

        Args:
            future_to_batch: Dictionary mapping futures to batch indices
            batch_iter: Iterator over remaining batches
            executor: Thread pool executor
            total_batches: Total number of batches

        Returns:
            The asset ingestion session response from the completed batch,
            or None if no batches completed (shouldn't happen)

        Raises:
            Exception: If the batch processing fails, all remaining batches are
                cancelled and the exception is re-raised

        """
        # 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_idx = future_to_batch[future]

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

                # Submit next batch if available
                self._submit_next_batch(
                    executor, batch_iter, future_to_batch, total_batches
                )

                return session_response

            except Exception as e:
                logger.error(
                    f"Batch {batch_idx} 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_iter: enumerate,
        future_to_batch: dict,
        total_batches: int,
    ):
        """Submit the next batch if one is available.

        This is called after a batch completes to maintain the desired level of
        parallelism. If no more batches are available, this does nothing.

        Args:
            executor: Thread pool executor
            batch_iter: Iterator over remaining batches
            future_to_batch: Dictionary mapping futures to batch indices
            total_batches: Total number of batches

        """
        try:
            batch_idx, batch_assets = next(batch_iter)
            future = self._submit_batch(
                executor, batch_idx, batch_assets, total_batches
            )
            future_to_batch[future] = batch_idx
        except StopIteration:
            pass

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

        This is called when a user cancels the upload or when a batch fails.
        All futures in the provided dictionary are cancelled.

        Args:
            future_to_batch: Dictionary of futures to cancel

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