#!/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    :   Datature Vi SDK assets uploader module.
"""

import os
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from threading import Lock

import httpx
import msgspec
from rich import print as rprint
from rich.progress import BarColumn, TaskID, TextColumn, TimeElapsedColumn
from vi.api.resources.datasets.assets import consts, responses
from vi.api.resources.datasets.assets.links import AssetIngestionSessionLinkParser
from vi.api.resources.datasets.assets.multipart import MultipartUploader
from vi.api.resources.datasets.assets.processor import ParallelBatchProcessor
from vi.api.resources.datasets.assets.types import (
    AssetBytesInput,
    AssetKind,
    AssetSource,
    AssetUploadFileSpec,
    AssetUploadSession,
    AssetUploadSpec,
    BatchConfig,
)
from vi.api.resources.datasets.utils.helper import (
    calculate_crc32c_and_detect_mime,
    calculate_crc32c_and_detect_mime_from_bytes,
    detect_mime_type,
)
from vi.api.resources.managers import ResourceUploader
from vi.client.errors import (
    ViInvalidParameterError,
    ViValidationError,
    extract_error_message_from_httpx_error,
)
from vi.client.http.retry import RetryExecutor, RetryHandler
from vi.client.validation import validate_id_param
from vi.logging import get_logger
from vi.utils.graceful_exit import GracefulExit, graceful_exit
from vi.utils.progress import ViProgress

logger = get_logger("assets.uploader")


class AssetUploadProgressTracker:
    """Thread-safe progress tracker for asset uploads."""

    def __init__(
        self, progress: ViProgress, task: TaskID, total_assets: int, total_bytes: int
    ):
        self.progress = progress
        self.task = task
        self.total_assets = total_assets
        self.total_bytes = total_bytes
        self.completed_count = 0
        self.completed_bytes = 0
        self._lock = Lock()

    @staticmethod
    def _format_bytes(bytes_value: int) -> str:
        """Format bytes to human-readable string."""
        if bytes_value >= 1024**3:  # >= 1 GB
            return f"{bytes_value / (1024**3):.2f} GB"
        if bytes_value >= 1024**2:  # >= 1 MB
            return f"{bytes_value / (1024**2):.2f} MB"
        if bytes_value >= 1024:  # >= 1 KB
            return f"{bytes_value / 1024:.2f} KB"
        return f"{bytes_value} B"

    def update_bytes(self, bytes_completed: int):
        """Update bytes progress without changing asset count (for multipart chunks)."""
        with self._lock:
            self.completed_bytes += bytes_completed
            uploaded_str = self._format_bytes(self.completed_bytes)
            total_str = self._format_bytes(self.total_bytes)
            self.progress.update(
                self.task,
                advance=bytes_completed,
                uploaded_str=uploaded_str,
                total_str=total_str,
            )

    def update_success(self, files_completed: int, bytes_completed: int):
        """Update progress for successful uploads."""
        with self._lock:
            self.completed_count += files_completed
            self.completed_bytes += bytes_completed
            uploaded_str = self._format_bytes(self.completed_bytes)
            total_str = self._format_bytes(self.total_bytes)
            self.progress.update(
                self.task,
                description=f"Uploading {self.completed_count} / {self.total_assets} assets...",
                advance=bytes_completed,
                uploaded_str=uploaded_str,
                total_str=total_str,
            )

    def update_error(self):
        """Update progress for failed uploads."""
        with self._lock:
            uploaded_str = self._format_bytes(self.completed_bytes)
            total_str = self._format_bytes(self.total_bytes)
            self.progress.update(
                self.task,
                description=f"Uploaded {self.completed_count} / {self.total_assets} assets... (error occurred)",
                uploaded_str=uploaded_str,
                total_str=total_str,
            )

    def finish_cancelled(self):
        """Update progress for cancelled uploads."""
        self.progress.update(self.task, description="✗ Upload cancelled")


class AssetUploader(ResourceUploader):
    """Uploader for assets with parallel batch processing support."""

    _link_parser: AssetIngestionSessionLinkParser | None = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Initialize retry executor with the requester's retry handler
        self._retry_executor = RetryExecutor(
            RetryHandler(self._requester._retry_config)
        )

    def upload(
        self,
        dataset_id: str,
        paths: Path | str | Sequence[Path | str] | None = None,
        assets_bytes: Sequence[AssetBytesInput] | None = None,
        failure_mode: str = "FailAfterOne",
        on_asset_overwritten: str = "RemoveLinkedResource",
        asset_source_class: str = consts.DEFAULT_ASSET_SOURCE_CLASS,
        asset_source_provider: str = consts.DEFAULT_ASSET_SOURCE_PROVIDER,
        asset_source_id: str = consts.DEFAULT_ASSET_SOURCE_ID,
        show_progress: bool = True,
        progress: ViProgress | None = None,
    ) -> list[responses.AssetIngestionSession]:
        """Upload assets to the dataset with automatic batching and parallelization.

        Args:
            dataset_id: ID of the dataset to upload to
            paths: File path(s) to upload - can be single path or list of paths
            assets_bytes: In-memory assets to upload as (filename, bytes) tuples
            failure_mode:
                How to handle failures ("FailAfterOne" or "FailAfterAll"), default is "FailAfterOne"
            on_asset_overwritten:
                Action when asset exists ("RemoveLinkedResource" or
                "KeepLinkedResource"), default is "RemoveLinkedResource"
            asset_source_class: Source class for assets
            asset_source_provider: Source provider for assets
            asset_source_id: Source ID for assets
            show_progress: Whether to display progress bars during upload
            progress: Optional existing ViProgress instance to use for progress display.
                If provided, progress bars will be added to this instance instead of
                creating a new one. This allows multiple operations to share the same
                progress display without flickering.

        Returns:
            List of asset ingestion session responses

        Raises:
            ViInvalidParameterError: If parameters are invalid
            ViValidationError: If no valid files found or both paths and assets_bytes are provided

        """
        # Validate and prepare
        self._validate_upload_params(dataset_id, failure_mode, on_asset_overwritten)
        asset_sources = self._collect_asset_sources(paths, assets_bytes)
        batches = self._split_into_batches(asset_sources, consts.MAX_UPLOAD_BATCH_SIZE)

        # Create configuration
        config = BatchConfig(
            dataset_id=dataset_id,
            failure_mode=failure_mode,
            on_asset_overwritten=on_asset_overwritten,
            asset_source_class=asset_source_class,
            asset_source_provider=asset_source_provider,
            asset_source_id=asset_source_id,
        )

        # Calculate total assets and estimate total bytes for ALL batches
        total_assets = len(asset_sources)
        total_bytes = self._estimate_total_bytes(asset_sources)

        # Warn user about preprocessing time if multiple batches
        if len(batches) > 1:
            rprint(
                "[yellow]⚠ Large upload detected:[/yellow] "
                f"Processing {total_assets} assets across {len(batches)} batches. "
                "Asset preprocessing may take some time. "
                "Upload progress will appear once preprocessing completes."
            )

        # Process batches
        with graceful_exit("Upload cancelled by user") as handler:
            # Determine progress instance to use
            if show_progress and progress is None:
                # Create and manage our own progress display
                with self._create_progress_display() as use_progress:
                    if len(batches) > 1 and consts.MAX_PARALLEL_BATCHES > 1:
                        return self._process_batches_parallel(
                            batches,
                            config,
                            use_progress,
                            handler,
                            total_assets,
                            total_bytes,
                        )
                    return self._process_batches_sequential(
                        batches,
                        config,
                        use_progress,
                        handler,
                        total_assets,
                        total_bytes,
                    )
            else:
                # Use provided progress instance or None
                use_progress = progress if show_progress else None
                if len(batches) > 1 and consts.MAX_PARALLEL_BATCHES > 1:
                    return self._process_batches_parallel(
                        batches,
                        config,
                        use_progress,
                        handler,
                        total_assets,
                        total_bytes,
                    )
                return self._process_batches_sequential(
                    batches, config, use_progress, handler, total_assets, total_bytes
                )

    def _validate_upload_params(
        self, dataset_id: str, failure_mode: str, on_asset_overwritten: str
    ):
        """Validate upload parameters.

        Ensures all required parameters are valid before starting the upload process.
        Raises appropriate exceptions for invalid values.

        Args:
            dataset_id: The dataset ID to validate
            failure_mode: The failure handling mode to validate
            on_asset_overwritten: The overwrite handling mode to validate

        Raises:
            ViInvalidParameterError: If any parameter is invalid

        """
        validate_id_param(dataset_id, "dataset_id")

        if failure_mode not in ["FailAfterOne", "FailAfterAll"]:
            raise ViInvalidParameterError(
                "failure_mode",
                f"Invalid failure_mode '{failure_mode}'. Must be 'FailAfterOne' or 'FailAfterAll'",
            )

        if on_asset_overwritten not in ["RemoveLinkedResource", "KeepLinkedResource"]:
            raise ViInvalidParameterError(
                "on_asset_overwritten",
                f"Invalid on_asset_overwritten '{on_asset_overwritten}'. "
                "Must be 'RemoveLinkedResource' or 'KeepLinkedResource'",
            )

    def _collect_asset_sources(
        self,
        paths: Path | str | Sequence[Path | str] | None,
        assets_bytes: Sequence[AssetBytesInput] | None,
    ) -> list[AssetSource]:
        """Collect all asset sources from the provided input.

        Parses the input which can be file paths, in-memory bytes, or both.
        Validates that only one input method is used.

        Args:
            paths: Single path, directory, or sequence of paths to collect files from
            assets_bytes: In-memory assets as (filename, bytes) tuples

        Returns:
            List of AssetSource objects ready for upload

        Raises:
            ViValidationError: If no valid assets are found or both inputs are provided

        """
        if paths is not None and assets_bytes is not None:
            raise ViValidationError(
                "Cannot specify both 'paths' and 'assets_bytes' parameters",
                suggestion="Use either 'paths' for file-based uploads or 'assets_bytes' for in-memory uploads",
            )

        if paths is None and assets_bytes is None:
            raise ViValidationError(
                "Must specify either 'paths' or 'assets_bytes' parameter",
                suggestion="Provide file paths using 'paths' or in-memory data using 'assets_bytes'",
            )

        asset_sources = []

        if paths is not None:
            file_paths = self._parse_asset_paths(paths)
            asset_sources.extend(
                [
                    AssetSource(filename=os.path.basename(path), file_path=path)
                    for path in file_paths
                ]
            )

        if assets_bytes is not None:
            asset_sources.extend(
                [
                    AssetSource(filename=filename, data=data)
                    for filename, data in assets_bytes
                ]
            )

        if not asset_sources:
            raise ViValidationError(
                "No valid assets found to upload",
                suggestion="Check that the provided paths exist and contain supported file formats, or that assets_bytes is not empty",
            )

        return asset_sources

    def _split_into_batches(
        self, asset_sources: list[AssetSource], batch_size: int
    ) -> list[list[AssetSource]]:
        """Split asset sources into batches of specified size.

        Divides the asset list into batches according to the specified batch size
        to respect API limits and optimize parallel processing.

        Args:
            asset_sources: List of asset sources to split
            batch_size: Maximum number of assets per batch

        Returns:
            List of batches, where each batch is a list of AssetSource objects

        """
        return [
            asset_sources[i : i + batch_size]
            for i in range(0, len(asset_sources), batch_size)
        ]

    def _estimate_total_bytes(self, asset_sources: list[AssetSource]) -> int:
        """Estimate total bytes for all asset sources.

        Args:
            asset_sources: List of asset sources

        Returns:
            Total estimated bytes across all assets

        """
        total = 0
        for asset_source in asset_sources:
            if asset_source.is_file() and asset_source.file_path:
                try:
                    total += os.path.getsize(asset_source.file_path)
                except OSError:
                    # If we can't get size, skip it (file might not exist)
                    pass
            elif asset_source.is_bytes() and asset_source.data:
                total += len(asset_source.data)
        return total

    def _create_progress_display(self) -> ViProgress:
        """Create the progress bar display for upload tracking.

        Returns:
            Configured ViProgress instance for displaying upload progress

        """
        return ViProgress(
            "[progress.description]{task.description}",
            BarColumn(),
            "[progress.percentage]{task.percentage:>3.2f}%",
            TextColumn("[cyan]{task.fields[uploaded_str]} / {task.fields[total_str]}"),
            TimeElapsedColumn(),
            transient=False,
        )

    def _process_batches_parallel(
        self,
        batches: list[list[AssetSource]],
        config: BatchConfig,
        progress: ViProgress | None,
        handler: GracefulExit,
        total_assets: int,
        total_bytes: int,
    ) -> list[responses.AssetIngestionSession]:
        """Process multiple batches in parallel with controlled concurrency.

        Uses the ParallelBatchProcessor to execute batches concurrently while
        respecting the MAX_PARALLEL_BATCHES limit. Implements dynamic batch
        submission to maximize throughput.

        Args:
            batches: List of asset source batches to process
            config: Batch configuration
            progress: Progress tracker
            handler: Graceful exit handler
            total_assets: Total number of assets across all batches
            total_bytes: Total bytes across all batches

        Returns:
            List of successful ingestion session responses

        """
        # Create master progress tracker for ALL batches
        master_tracker = None
        if progress:
            total_str = AssetUploadProgressTracker._format_bytes(total_bytes)
            upload_task = progress.add_task(
                f"Uploading 0 / {total_assets} assets...",
                total=total_bytes,
                uploaded_str="0 B",
                total_str=total_str,
            )
            master_tracker = AssetUploadProgressTracker(
                progress, upload_task, total_assets, total_bytes
            )

        processor = ParallelBatchProcessor(
            self, config, progress, handler, master_tracker
        )
        max_parallel = min(consts.MAX_PARALLEL_BATCHES, len(batches))

        try:
            return processor.process_batches(batches, max_parallel)
        finally:
            # Clean up master tracker
            if progress and master_tracker:
                progress.stop()
                # Only show success message if not cancelled
                if not handler.exit_now:
                    uploaded_str = AssetUploadProgressTracker._format_bytes(
                        master_tracker.completed_bytes
                    )
                    print(
                        f"✓ Uploaded {master_tracker.completed_count} / {total_assets} "
                        f"assets ({uploaded_str})",
                        flush=True,
                    )
                progress.remove_task(master_tracker.task)

    def _process_batches_sequential(
        self,
        batches: list[list[AssetSource]],
        config: BatchConfig,
        progress: ViProgress | None,
        handler: GracefulExit,
        total_assets: int,
        total_bytes: int,
    ) -> list[responses.AssetIngestionSession]:
        """Process batches one at a time in sequence.

        Used when there's only one batch or when parallel processing is disabled.
        Processes each batch completely before moving to the next.

        Args:
            batches: List of asset source batches to process
            config: Batch configuration
            progress: Progress tracker
            handler: Graceful exit handler
            total_assets: Total number of assets across all batches
            total_bytes: Total bytes across all batches

        Returns:
            List of successful ingestion session responses

        """
        # Create master progress tracker for ALL batches
        master_tracker = None
        if progress:
            total_str = AssetUploadProgressTracker._format_bytes(total_bytes)
            upload_task = progress.add_task(
                f"Uploading 0 / {total_assets} assets...",
                total=total_bytes,
                uploaded_str="0 B",
                total_str=total_str,
            )
            master_tracker = AssetUploadProgressTracker(
                progress, upload_task, total_assets, total_bytes
            )

        session_responses = []
        total_batches = len(batches)

        try:
            for batch_idx, batch_assets in enumerate(batches, start=1):
                session_response = self.process_single_batch(
                    config=config,
                    batch_assets=batch_assets,
                    batch_idx=batch_idx,
                    total_batches=total_batches,
                    progress=progress,
                    handler=handler,
                    master_tracker=master_tracker,
                )

                if session_response is None:
                    break  # Cancelled

                session_responses.append(session_response)

            return session_responses
        finally:
            # Clean up master tracker
            if progress and master_tracker:
                progress.stop()
                # Only show success message if not cancelled
                if not handler.exit_now:
                    uploaded_str = AssetUploadProgressTracker._format_bytes(
                        master_tracker.completed_bytes
                    )
                    print(
                        f"✓ Uploaded {master_tracker.completed_count} / {total_assets} "
                        f"assets ({uploaded_str})",
                        flush=True,
                    )
                progress.remove_task(master_tracker.task)

    def process_single_batch(
        self,
        config: BatchConfig,
        batch_assets: list[AssetSource],
        batch_idx: int,
        total_batches: int,
        progress: ViProgress | None,
        handler: GracefulExit,
        master_tracker: AssetUploadProgressTracker | None = None,
    ) -> responses.AssetIngestionSession | None:
        """Process a single batch (preparation + upload).

        Args:
            config: Batch configuration
            batch_assets: List of asset sources in this batch
            batch_idx: Index of this batch (1-based)
            total_batches: Total number of batches
            progress: Progress tracker
            handler: Graceful exit handler
            master_tracker: Optional master progress tracker for all batches

        Returns:
            Asset ingestion session response or None if cancelled

        """
        # Early cancellation check
        if handler.exit_now:
            return None

        batch_info = (
            f" (batch {batch_idx}/{total_batches})" if total_batches > 1 else ""
        )

        # Step 1: Prepare assets
        try:
            session_response = self._prepare_batch(config, batch_assets, handler)

            if session_response is None:
                return None  # Cancelled

        except KeyboardInterrupt:
            # KeyboardInterrupt from prepare phase - session already aborted if it was created
            logger.debug(f"KeyboardInterrupt during batch preparation {batch_info}")
            raise

        # Step 2: Upload assets
        try:
            # Ensure we have assets to upload
            if not session_response.status.assets:
                logger.warning(f"No assets to upload in batch{batch_info}")
                return session_response

            self._upload_batch(
                batch_assets,
                session_response.status.assets,
                session_response.ingestion_id,
                handler,
                master_tracker,
            )

            # Check if cancelled during upload
            if handler.exit_now:
                logger.info(
                    f"Upload cancelled, aborting ingestion session {session_response.asset_ingestion_session_id}"
                )
                self._abort_ingestion_session(
                    session_response.asset_ingestion_session_id
                )
                return None

        except KeyboardInterrupt:
            # Abort ingestion session on keyboard interrupt to prevent resource leak
            logger.info(
                f"KeyboardInterrupt during batch upload, aborting ingestion session "
                f"{session_response.asset_ingestion_session_id}"
            )
            self._abort_ingestion_session(session_response.asset_ingestion_session_id)
            raise
        except Exception as e:
            # Abort ingestion session on error
            logger.error(
                f"Error during batch upload, aborting ingestion session "
                f"{session_response.asset_ingestion_session_id}: {e}"
            )
            self._abort_ingestion_session(session_response.asset_ingestion_session_id)
            raise

        return session_response

    def _prepare_batch(
        self,
        config: BatchConfig,
        batch_assets: list[AssetSource],
        handler: GracefulExit,
    ) -> responses.AssetIngestionSession | None:
        """Prepare a batch for upload by generating metadata and creating a session.

        This performs two main steps:
        1. Generate file specifications (MIME type, CRC32C, etc.) for each asset
        2. Create an ingestion session via API with the generated specs

        Args:
            config: Batch configuration with upload parameters
            batch_assets: List of asset sources in this batch
            handler: Graceful exit handler

        Returns:
            Asset ingestion session response, or None if cancelled

        Raises:
            ValueError: If the API returns an invalid response

        """
        session = None
        try:
            # Generate file specs (silently)
            assets = self._generate_file_specs(batch_assets, handler)

            if assets is None:
                return None  # Cancelled

            # Check for cancellation before making API call
            if handler.exit_now:
                return None

            # Create ingestion session (silently)
            session = self._create_ingestion_session(config, assets, handler)

            return session

        except KeyboardInterrupt:
            # If session was created but then KeyboardInterrupt occurred,
            # abort the session to prevent resource leak
            if session is not None:
                logger.info(
                    f"Aborting ingestion session {session.asset_ingestion_session_id} due to KeyboardInterrupt"
                )
                self._abort_ingestion_session(session.asset_ingestion_session_id)
            # Re-raise KeyboardInterrupt to propagate cancellation signal
            raise

    def _generate_file_specs(
        self,
        batch_assets: list[AssetSource],
        handler: GracefulExit,
    ) -> list[AssetUploadFileSpec] | None:
        """Generate file specifications for all assets in the batch.

        For each asset, generates metadata including:
        - File size and name
        - MIME type
        - CRC32C checksum
        - Asset kind (Image/Video)

        This is the most time-consuming part of preparation as it reads
        each asset to calculate checksums. Uses parallel processing for
        batches larger than FILE_SPEC_PARALLEL_THRESHOLD.

        Args:
            batch_assets: List of asset sources to generate specs for
            handler: Graceful exit handler

        Returns:
            List of file specifications, or None if cancelled

        """
        # Use parallel processing for larger batches to maximize throughput
        if len(batch_assets) >= consts.FILE_SPEC_PARALLEL_THRESHOLD:
            return self._generate_file_specs_parallel(batch_assets, handler)

        # Sequential processing for small batches (less overhead)
        return self._generate_file_specs_sequential(batch_assets, handler)

    def _generate_file_specs_sequential(
        self,
        batch_assets: list[AssetSource],
        handler: GracefulExit,
    ) -> list[AssetUploadFileSpec] | None:
        """Generate file specifications sequentially (used for small batches).

        Args:
            batch_assets: List of asset sources to generate specs for
            handler: Graceful exit handler

        Returns:
            List of file specifications, or None if cancelled

        """
        assets = []

        for asset_source in batch_assets:
            if handler.exit_now:
                return None  # Cancelled

            file_spec = self._generate_asset_file_spec(asset_source)
            if file_spec is not None:
                assets.append(file_spec)

        return assets

    def _generate_file_specs_parallel(
        self,
        batch_assets: list[AssetSource],
        handler: GracefulExit,
    ) -> list[AssetUploadFileSpec] | None:
        """Generate file specifications in parallel (used for large batches).

        Uses a ThreadPoolExecutor to process multiple assets concurrently,
        significantly improving throughput for I/O-bound operations like
        reading files for CRC32C calculation.

        Args:
            batch_assets: List of asset sources to generate specs for
            handler: Graceful exit handler

        Returns:
            List of file specifications, or None if cancelled

        """
        assets: list[AssetUploadFileSpec] = []
        lock = Lock()

        def process_asset(asset_source: AssetSource) -> AssetUploadFileSpec | None:
            """Process a single asset.

            Note: KeyboardInterrupt is allowed to propagate from file operations
            to ensure immediate cancellation when user presses Ctrl+C.
            """
            # Check if cancellation was requested
            if handler.exit_now:
                return None

            file_spec = self._generate_asset_file_spec(asset_source)
            return file_spec

        executor = ThreadPoolExecutor(max_workers=consts.MAX_PARALLEL_FILE_SPEC_WORKERS)
        try:
            # Submit all assets for processing
            futures = {
                executor.submit(process_asset, asset): asset for asset in batch_assets
            }

            # Collect results as they complete with periodic cancellation checks
            remaining_futures = set(futures.keys())
            while remaining_futures:
                if handler.exit_now:
                    # Cancel remaining futures
                    for f in remaining_futures:
                        f.cancel()
                    return None

                try:
                    # Check for completed futures with short timeout
                    for future in as_completed(remaining_futures, timeout=0.1):
                        remaining_futures.discard(future)
                        try:
                            file_spec = future.result()
                            if file_spec is not None:
                                # Thread-safe append
                                with lock:
                                    assets.append(file_spec)
                        except Exception as e:
                            asset_source = futures[future]
                            logger.error(
                                f"Error processing asset {asset_source.filename}: {e}"
                            )
                            # Continue processing other assets
                        break  # Process one future at a time to check handler.exit_now
                except TimeoutError:
                    # No futures completed in timeout period, loop again to check exit_now
                    continue

        finally:
            # Ensure executor shuts down immediately without waiting
            executor.shutdown(wait=False, cancel_futures=True)

        return assets

    def _create_ingestion_session(
        self,
        config: BatchConfig,
        assets: list[AssetUploadFileSpec],
        handler: GracefulExit | None = None,
    ) -> responses.AssetIngestionSession:
        """Create an asset ingestion session via API.

        Sends the file specifications to the API to create an ingestion session.
        The API responds with signed URLs for uploading each asset.

        Args:
            config: Batch configuration
            assets: List of file specifications
            handler: Optional graceful exit handler for cancellation support

        Returns:
            Asset ingestion session with upload URLs

        Raises:
            ValueError: If the API response is invalid
            KeyboardInterrupt: If cancellation was requested before making the request

        """
        # Final cancellation check before making the HTTP request
        if handler and handler.exit_now:
            raise KeyboardInterrupt(
                "Upload cancelled before creating ingestion session"
            )

        self._link_parser = AssetIngestionSessionLinkParser(
            self._auth.organization_id,
            config.dataset_id,
            config.asset_source_class,
            config.asset_source_provider,
            config.asset_source_id,
        )

        upload_session = AssetUploadSession(
            spec=AssetUploadSpec(
                assets=assets,
                failure_mode=config.failure_mode,
                on_asset_overwritten=config.on_asset_overwritten,
            )
        )

        response = self._requester.post(
            self._link_parser(),
            response_type=responses.AssetIngestionSession,
            json_data=msgspec.to_builtins(upload_session, str_keys=True),
        )

        if not isinstance(response, responses.AssetIngestionSession):
            raise ValueError(f"Invalid response {response} with type {type(response)}")

        return response

    def _upload_batch(
        self,
        batch_assets: list[AssetSource],
        assets_for_upload: list[responses.AssetForUpload],
        ingestion_id: str,
        handler: GracefulExit,
        master_tracker: AssetUploadProgressTracker | None,
    ):
        """Upload all assets in a batch using their signed URLs.

        Performs the actual file uploads using the signed URLs provided by the
        ingestion session. Uses parallel threads for efficient upload.

        Args:
            batch_assets: List of asset sources to upload
            assets_for_upload: List of asset upload info with signed URLs
            ingestion_id: The ingestion ID for multipart uploads
            handler: Graceful exit handler
            master_tracker: Master progress tracker (None if progress is disabled)

        """
        # Simply use the master tracker (or None if progress disabled)
        self._upload_assets_parallel(
            batch_assets,
            assets_for_upload,
            ingestion_id,
            master_tracker,
            handler,
        )

    def _upload_assets_parallel(
        self,
        batch_assets: list[AssetSource],
        assets_for_upload: list[responses.AssetForUpload],
        ingestion_id: str,
        progress_tracker: AssetUploadProgressTracker | None,
        handler: GracefulExit,
    ):
        """Upload assets in parallel using thread pool."""
        # Match asset sources to assets
        matched_assets = self._match_assets_to_upload_info(
            batch_assets, assets_for_upload
        )

        if not matched_assets:
            return

        # Upload in parallel
        executor = ThreadPoolExecutor(max_workers=self._max_workers)
        try:
            futures = {
                executor.submit(
                    self._upload_single_asset_with_progress,
                    asset_source,
                    asset.upload,
                    ingestion_id,
                    asset.metadata.size,
                    progress_tracker,
                    handler,
                ): asset_source
                for asset_source, asset in matched_assets
            }

            # Wait for completion with periodic cancellation checks
            # When cancellation is detected, we still wait for all workers to finish
            # so they can properly abort multipart uploads and clean up resources
            remaining_futures = set(futures.keys())
            cancellation_detected = False

            while remaining_futures:
                if handler.exit_now and not cancellation_detected:
                    cancellation_detected = True
                    if progress_tracker:
                        progress_tracker.finish_cancelled()
                    logger.info(
                        f"Cancellation detected, waiting for {len(remaining_futures)} "
                        f"upload workers to finish cleanup"
                    )
                    # Cancel futures that haven't started yet
                    for f in remaining_futures:
                        f.cancel()

                try:
                    # Check for completed futures with short timeout
                    for future in as_completed(remaining_futures, timeout=0.1):
                        remaining_futures.discard(future)
                        try:
                            future.result()
                        except httpx.HTTPStatusError as e:
                            error_msg = extract_error_message_from_httpx_error(e)
                            rprint(f"[red]Asset upload failed:[/red]\n{error_msg}")
                        except (
                            OSError,
                            httpx.TimeoutException,
                            httpx.NetworkError,
                        ) as e:
                            rprint(f"[red]Asset upload failed: {e}[/red]")
                        break  # Process one future at a time to check handler.exit_now
                except TimeoutError:
                    # No futures completed in timeout period, loop again to check exit_now
                    continue

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

    def _match_assets_to_upload_info(
        self,
        batch_assets: list[AssetSource],
        assets: list[responses.AssetForUpload],
    ) -> list[tuple[AssetSource, responses.AssetForUpload]]:
        """Match asset sources to their corresponding asset upload information.

        The API returns asset upload info with filenames. This method matches
        those filenames back to the original asset sources for uploading.

        Args:
            batch_assets: List of original asset sources
            assets: List of asset upload info from the API

        Returns:
            List of (asset_source, asset_upload_info) tuples ready for upload

        """
        matched = []
        for asset in assets:
            for asset_source in batch_assets:
                if asset_source.filename == asset.metadata.filename:
                    matched.append((asset_source, asset))
                    break
        return matched

    def _upload_single_asset_with_progress(
        self,
        asset_source: AssetSource,
        signed_url_info: responses.Contents,
        ingestion_id: str,
        asset_size: int,
        progress_tracker: AssetUploadProgressTracker | None,
        handler: GracefulExit,
    ) -> bool:
        """Upload a single asset with progress tracking and retry logic.

        Args:
            asset_source: Asset source to upload
            signed_url_info: Signed URL and headers for upload
            ingestion_id: The ingestion ID for multipart uploads
            asset_size: Size of the asset in bytes
            progress_tracker: Optional progress tracker for UI updates
            handler: Graceful exit handler for cancellation

        Returns:
            True if upload successful, False if cancelled

        Raises:
            Exception: If upload fails after retries

        """
        if handler.exit_now:
            return False

        try:
            was_multipart = self._retry_executor.execute(
                operation=lambda: self._upload_single_asset(
                    asset_source,
                    signed_url_info,
                    ingestion_id,
                    handler,
                    progress_tracker,
                ),
                operation_name=f"Upload {asset_source.filename}",
                max_retries=3,
                logger=logger,
                context_info={"file": asset_source.filename},
            )
            # Only update progress if not cancelled
            if progress_tracker and not handler.exit_now:
                # For multipart, bytes were already counted via update_bytes()
                # Only count asset and bytes for non-multipart uploads
                if was_multipart:
                    progress_tracker.update_success(1, 0)  # Just count the asset
                else:
                    progress_tracker.update_success(
                        1, asset_size
                    )  # Count asset + bytes
            return True
        except httpx.HTTPStatusError as e:
            error_msg = extract_error_message_from_httpx_error(e)
            rprint(f"[red]Failed to upload {asset_source.filename}:[/red]\n{error_msg}")
            if progress_tracker:
                progress_tracker.update_error()
            raise
        except (OSError, httpx.TimeoutException, httpx.NetworkError) as e:
            rprint(f"[red]Failed to upload {asset_source.filename}: {e}[/red]")
            if progress_tracker:
                progress_tracker.update_error()
            raise

    def _upload_single_asset(
        self,
        asset_source: AssetSource,
        signed_url_info: responses.Contents,
        ingestion_id: str,
        handler: GracefulExit,
        progress_tracker: AssetUploadProgressTracker | None = None,
    ) -> bool:
        """Upload an asset using signed URL or multipart upload for large files.

        This is the actual upload operation without retry logic.
        Retries are handled by the caller using RetryExecutor.

        Args:
            asset_source: Asset source to upload (file path or bytes)
            signed_url_info: Signed URL and headers for upload
            ingestion_id: The ingestion ID for multipart uploads
            handler: Graceful exit handler for cancellation
            progress_tracker: Optional progress tracker for bytes tracking

        Returns:
            True if multipart upload was used, False otherwise
            ingestion_id: The ingestion ID for multipart uploads
            handler: Graceful exit handler for cancellation
            progress_tracker: Optional progress tracker for UI updates

        Raises:
            httpx.HTTPStatusError: If upload fails
            OSError: If file cannot be read

        """
        if asset_source.is_file():
            file_size = os.path.getsize(asset_source.file_path)

            # Use multipart for large files
            if file_size >= consts.MULTIPART_THRESHOLD_BYTES:
                uploader = MultipartUploader(
                    uploader=self,
                    file_path=asset_source.file_path,
                    filename=asset_source.filename,
                    max_workers=consts.MULTIPART_UPLOAD_WORKERS,
                )

                uploader.upload(
                    ingestion_id=ingestion_id,
                    handler=handler,
                    progress_tracker=progress_tracker,
                )

                return True  # Was multipart

            # Standard single-request upload
            with open(asset_source.file_path, "rb") as f:
                response = self._get_upload_http_client().put(
                    signed_url_info.url,
                    content=f,
                    headers=signed_url_info.headers,
                )
                response.raise_for_status()
                return False  # Was not multipart

        elif asset_source.is_bytes():
            # Upload from in-memory bytes
            response = self._get_upload_http_client().put(
                signed_url_info.url,
                content=asset_source.data,
                headers=signed_url_info.headers,
            )
            response.raise_for_status()
            return False  # Was not multipart
        else:
            raise ValueError(
                f"Invalid AssetSource: neither file_path nor data is set for {asset_source.filename}"
            )

    def _parse_asset_paths(self, paths: Path | str | Sequence[Path | str]) -> list[str]:
        """Parse and expand asset paths into a list of file paths.

        Handles various input formats:
        - Single file path -> returns list with one file
        - Directory path -> returns all supported files in directory (recursive)
        - List of paths -> processes each and combines results

        Args:
            paths: Single path, directory, or sequence of paths

        Returns:
            List of absolute file paths

        Raises:
            ViInvalidParameterError: If a path doesn't exist or isn't a file/directory

        """
        path_list = [paths] if isinstance(paths, (str, Path)) else list(paths)
        file_paths = []

        for path_item in path_list:
            path = Path(path_item).expanduser().resolve()

            if path.is_file():
                file_paths.append(str(path))
            elif path.is_dir():
                file_paths.extend(self._scan_directory(path))
            else:
                raise ViInvalidParameterError(
                    "paths",
                    f"Invalid path: {path_item}. Path must be a file or directory",
                )

        return file_paths

    def _scan_directory(self, directory: Path) -> list[str]:
        """Scan directory recursively for supported asset files.

        Searches the directory and all subdirectories for files with supported
        extensions (images, videos, etc.).

        Args:
            directory: Path to directory to scan

        Returns:
            List of absolute paths to supported files found

        """
        return [
            str(file_path)
            for file_path in directory.glob("**/*")
            if file_path.is_file()
            and file_path.suffix.lower() in consts.SUPPORTED_ASSET_FILE_EXTENSIONS
        ]

    def _generate_asset_file_spec(
        self, asset_source: AssetSource
    ) -> AssetUploadFileSpec | None:
        """Generate complete file specification for a single asset.

        Analyzes the asset to extract all required metadata:
        - File size and name
        - MIME type (from extension or file content)
        - CRC32C checksum (for integrity verification)
        - Asset kind (Image/Video)

        Args:
            asset_source: Asset source (file path or bytes)

        Returns:
            Complete file specification ready for upload, or None if the asset
            should be skipped (e.g., MIME type validation fails)

        Raises:
            ValueError: If MIME type is valid but asset kind is unsupported

        """
        if asset_source.is_file():
            return self._generate_asset_file_spec_from_path(asset_source)
        if asset_source.is_bytes():
            return self._generate_asset_file_spec_from_bytes(asset_source)

        logger.warning(
            f"Skipping {asset_source.filename}: Invalid AssetSource (neither file_path nor data is set)"
        )
        return None

    def _generate_asset_file_spec_from_path(
        self, asset_source: AssetSource
    ) -> AssetUploadFileSpec | None:
        """Generate file spec from a file path.

        Args:
            asset_source: Asset source with file_path set

        Returns:
            Complete file specification or None if validation fails

        """
        asset_path = Path(asset_source.file_path).expanduser().resolve()

        # Get basic file metadata
        size = os.stat(asset_path).st_size
        filename = asset_source.filename
        ext = asset_path.suffix.lower()

        # Get expected MIME type from extension
        mime_from_ext = consts.ASSET_FILE_EXTENSION_TO_MIME_TYPE_MAP.get(ext)
        if not mime_from_ext:
            logger.warning(f"Skipping {asset_path}: Unknown file extension '{ext}'")
            return None

        # For files larger than multipart threshold, skip CRC32C calculation
        if size > consts.MULTIPART_THRESHOLD_BYTES:
            # Only detect MIME type without calculating checksum
            mime_from_content = detect_mime_type(asset_path)
            crc32c_value = None
        else:
            # Calculate CRC32C and detect MIME type
            crc32c_value, mime_from_content = calculate_crc32c_and_detect_mime(
                asset_path
            )

        # Validate MIME type from file content
        if not mime_from_content:
            logger.warning(
                f"Skipping {asset_path}: Unable to detect MIME type from file content"
            )
            return None

        # Verify that extension-based MIME matches content-detected MIME
        mime_base_ext = mime_from_ext.split("/")[0]
        mime_base_content = mime_from_content.split("/")[0]

        if mime_base_ext != mime_base_content:
            logger.warning(
                f"Skipping {asset_path}: MIME type mismatch - "
                f"extension suggests '{mime_from_ext}' but content is '{mime_from_content}'"
            )
            return None

        # Determine asset kind
        kind = self._determine_asset_kind(mime_from_ext)

        return AssetUploadFileSpec(
            filename=filename,
            mime=mime_from_ext,
            size=size,
            crc32c=crc32c_value,
            kind=kind,
        )

    def _generate_asset_file_spec_from_bytes(
        self, asset_source: AssetSource
    ) -> AssetUploadFileSpec | None:
        """Generate file spec from in-memory bytes.

        Args:
            asset_source: Asset source with data set

        Returns:
            Complete file specification or None if validation fails

        """
        filename = asset_source.filename
        data = asset_source.data
        size = len(data)

        # Get file extension from filename
        ext = Path(filename).suffix.lower()

        # Get expected MIME type from extension
        mime_from_ext = consts.ASSET_FILE_EXTENSION_TO_MIME_TYPE_MAP.get(ext)

        # Calculate CRC32C and detect MIME type from bytes
        crc32c_value, mime_from_content = calculate_crc32c_and_detect_mime_from_bytes(
            data
        )

        # Validate MIME type from content
        if not mime_from_content:
            logger.warning(
                f"Skipping {filename}: Unable to detect MIME type from content"
            )
            return None

        # If we have extension-based MIME, verify it matches content
        if mime_from_ext:
            mime_base_ext = mime_from_ext.split("/")[0]
            mime_base_content = mime_from_content.split("/")[0]

            if mime_base_ext != mime_base_content:
                logger.warning(
                    f"Skipping {filename}: MIME type mismatch - "
                    f"extension suggests '{mime_from_ext}' but content is '{mime_from_content}'"
                )
                return None

            # Use extension-based MIME as it's more specific
            final_mime = mime_from_ext
        else:
            # No extension or unknown extension, use content-detected MIME
            final_mime = mime_from_content

        # Determine asset kind and return spec
        kind = self._determine_asset_kind(final_mime)

        return AssetUploadFileSpec(
            filename=filename,
            mime=final_mime,
            size=size,
            crc32c=crc32c_value,
            kind=kind,
        )

    def _determine_asset_kind(self, mime_type: str) -> AssetKind:
        """Determine the asset kind from its MIME type.

        Maps MIME types to asset kinds:
        - image/* -> AssetKind.IMAGE
        - video/* -> AssetKind.VIDEO

        Args:
            mime_type: The MIME type string

        Returns:
            AssetKind enum value

        Raises:
            ValueError: If the MIME type is not supported

        """
        if mime_type.startswith("image/"):
            return AssetKind.IMAGE
        if mime_type.startswith("video/"):
            return AssetKind.VIDEO
        raise ValueError(f"Unsupported MIME type: {mime_type}")

    def get_asset_ingestion_session(
        self,
        dataset_id: str,
        asset_ingestion_session_id: str,
        asset_source_class: str = consts.DEFAULT_ASSET_SOURCE_CLASS,
        asset_source_provider: str = consts.DEFAULT_ASSET_SOURCE_PROVIDER,
        asset_source_id: str = consts.DEFAULT_ASSET_SOURCE_ID,
    ) -> responses.AssetIngestionSession:
        """Get an asset ingestion session by ID.

        Args:
            dataset_id: ID of the dataset to get the ingestion session for
            asset_ingestion_session_id: ID of the ingestion session to get
            asset_source_class: Class of the asset source
            asset_source_provider: Provider of the asset source
            asset_source_id: ID of the asset source

        Returns:
            Asset ingestion session response

        Raises:
            ValueError: If the response is invalid

        """
        self._link_parser = AssetIngestionSessionLinkParser(
            self._auth.organization_id,
            dataset_id,
            asset_source_class,
            asset_source_provider,
            asset_source_id,
        )

        response = self._requester.get(
            self._link_parser(asset_ingestion_session_id),
            response_type=responses.AssetIngestionSession,
        )

        if not isinstance(response, responses.AssetIngestionSession):
            raise ValueError(f"Invalid response {response} with type {type(response)}")

        return response

    def _abort_ingestion_session(self, asset_ingestion_session_id: str) -> None:
        """Abort an ingestion session to clean up server-side resources.

        This prevents resource leaks when upload fails or is cancelled.
        Failures are logged but not raised to avoid masking the original error.

        Args:
            asset_ingestion_session_id: Asset ingestion session ID to abort.

        """
        if not self._link_parser:
            logger.warning(
                f"Cannot abort ingestion session {asset_ingestion_session_id}: link parser not initialized"
            )
            return

        try:
            self._requester.post(
                f"{self._link_parser(asset_ingestion_session_id)}:cancel"
            )
            logger.info(f"Aborted ingestion session {asset_ingestion_session_id}")
        except Exception as e:
            # Log but don't raise - we don't want to mask the original error
            logger.warning(
                f"Failed to abort ingestion session {asset_ingestion_session_id}: {e}"
            )
