#!/usr/bin/env python

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

@File    :   progress.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Progress tracking for HuggingFace dataset uploads.
"""

from __future__ import annotations

from threading import Lock

from rich.progress import TaskID
from vi.utils.progress import ViProgress


class HFBatchProgressTracker:
    """Thread-safe progress tracker for HF batch uploads.

    Tracks both aggregate progress and individual batch progress tasks.
    When multiple batches are processed in parallel, creates separate
    progress tasks for each batch to show concurrent upload status.
    """

    def __init__(
        self,
        progress: ViProgress | None,
        total_images: int,
        show_batch_details: bool = False,
    ):
        """Initialize progress tracker.

        Args:
            progress: Progress display instance
            total_images: Total number of images to upload
            show_batch_details: If True, create individual progress tasks for each batch

        """
        self.progress = progress
        self.total_images = total_images
        self.completed_images = 0
        self.show_batch_details = show_batch_details
        self._lock = Lock()

        # Map batch_num -> TaskID for individual batch progress tasks
        self._batch_tasks: dict[int, TaskID] = {}

        # Overall aggregate task
        self.aggregate_task: TaskID | None = None
        if progress and not show_batch_details:
            # Only show aggregate task if not showing batch details
            self.aggregate_task = progress.add_task(
                "Uploading saved images to platform...",
                total=total_images,
            )

    def start_batch(self, batch_num: int, batch_size: int):
        """Create progress task when a batch starts uploading.

        Args:
            batch_num: Batch number (0-indexed)
            batch_size: Number of images in this batch

        """
        if not self.progress or not self.show_batch_details:
            return

        with self._lock:
            task = self.progress.add_task(
                f"Batch {batch_num + 1}: Uploading {batch_size} images",
                total=batch_size,
            )
            self._batch_tasks[batch_num] = task

    def finish_batch(self, batch_num: int, images_count: int):
        """Update progress when a batch completes.

        Args:
            batch_num: Batch number that completed
            images_count: Number of images processed in this batch

        """
        with self._lock:
            self.completed_images += images_count

            if self.progress:
                # Update or complete individual batch task
                if self.show_batch_details and batch_num in self._batch_tasks:
                    batch_task = self._batch_tasks[batch_num]
                    self.progress.update(
                        batch_task,
                        description=f"Batch {batch_num + 1}: ✓ Uploaded {images_count} images",
                        completed=images_count,
                    )
                    # Remove task after a brief delay
                    self.progress.remove_task(batch_task)
                    del self._batch_tasks[batch_num]

                # Update aggregate task
                elif self.aggregate_task is not None:
                    self.progress.update(
                        self.aggregate_task,
                        description=f"Uploaded {self.completed_images} / {self.total_images} images to platform",
                        advance=images_count,
                    )

    def finish_cancelled(self):
        """Update progress for cancelled uploads."""
        if not self.progress:
            return

        with self._lock:
            # Cancel all batch tasks
            for batch_num, task in self._batch_tasks.items():
                self.progress.update(
                    task, description=f"Batch {batch_num + 1}: ✗ Cancelled"
                )
                self.progress.remove_task(task)
            self._batch_tasks.clear()

            # Update aggregate task
            if self.aggregate_task is not None:
                self.progress.update(
                    self.aggregate_task, description="✗ Upload cancelled"
                )
