"""Argus CLI - Vision AI dataset toolkit."""

import hashlib
from pathlib import Path
from typing import Annotated

import cv2
import numpy as np
import typer
from rich.console import Console
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
)
from rich.table import Table

from argus.core import COCODataset, Dataset, MaskDataset, YOLODataset
from argus.core.base import DatasetFormat, TaskType
from argus.core.convert import convert_mask_to_yolo_seg
from argus.core.split import (
    is_coco_unsplit,
    parse_ratio,
    split_coco_dataset,
    split_yolo_dataset,
)

console = Console()

app = typer.Typer(
    name="argus",
    help="Vision AI dataset toolkit for working with YOLO and COCO datasets.",
    no_args_is_help=True,
)


@app.callback()
def callback() -> None:
    """Vision AI dataset toolkit for working with YOLO and COCO datasets."""
    pass


@app.command(name="list")
def list_datasets(
    path: Annotated[
        Path,
        typer.Option(
            "--path",
            "-p",
            help="Root path to search for datasets.",
        ),
    ] = Path("."),
    max_depth: Annotated[
        int,
        typer.Option(
            "--max-depth",
            "-d",
            help="Maximum directory depth to search.",
            min=1,
            max=10,
        ),
    ] = 3,
) -> None:
    """List all detected datasets in the specified path.

    Searches for YOLO and COCO format datasets within the given directory,
    up to the specified maximum depth.
    """
    # Resolve path and validate
    path = path.resolve()
    if not path.exists():
        console.print(f"[red]Error: Path does not exist: {path}[/red]")
        raise typer.Exit(1)
    if not path.is_dir():
        console.print(f"[red]Error: Path is not a directory: {path}[/red]")
        raise typer.Exit(1)

    datasets = _discover_datasets(path, max_depth)

    if not datasets:
        console.print(f"[yellow]No datasets found in {path}[/yellow]")
        return

    # Create and populate table
    table = Table(title=f"Datasets found in {path}")
    table.add_column("Path", style="cyan", no_wrap=True)
    table.add_column("Format", style="green")
    table.add_column("Task", style="magenta")
    table.add_column("Classes", justify="right", style="yellow")
    table.add_column("Splits", style="blue")

    for dataset in datasets:
        summary = dataset.summary()
        table.add_row(
            summary["path"],
            summary["format"],
            summary["task"],
            str(summary["classes"]),
            summary["splits"],
        )

    console.print(table)
    console.print(f"\n[green]Found {len(datasets)} dataset(s)[/green]")


@app.command(name="stats")
def stats(
    dataset_path: Annotated[
        Path,
        typer.Option(
            "--dataset-path",
            "-d",
            help="Path to the dataset root directory.",
        ),
    ] = Path("."),
) -> None:
    """Show instance statistics for a dataset.

    Displays the number of annotation instances per class, per split.
    The path should point to a dataset root containing data.yaml (YOLO)
    or an annotations/ folder (COCO).
    """
    # Resolve path and validate
    dataset_path = dataset_path.resolve()
    if not dataset_path.exists():
        console.print(f"[red]Error: Path does not exist: {dataset_path}[/red]")
        raise typer.Exit(1)
    if not dataset_path.is_dir():
        console.print(f"[red]Error: Path is not a directory: {dataset_path}[/red]")
        raise typer.Exit(1)

    # Detect dataset
    dataset = _detect_dataset(dataset_path)
    if not dataset:
        console.print(
            f"[red]Error: No dataset found at {dataset_path}[/red]\n"
            "[yellow]Ensure the path points to a dataset root containing "
            "data.yaml (YOLO), annotations/ folder (COCO), or "
            "images/ + masks/ directories (Mask).[/yellow]"
        )
        raise typer.Exit(1)

    # Handle mask datasets with pixel statistics
    if dataset.format == DatasetFormat.MASK:
        assert isinstance(dataset, MaskDataset)
        _show_mask_stats(dataset, dataset_path)
        return

    # Get instance counts with progress indicator
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
        transient=True,
    ) as progress:
        progress.add_task("Analyzing dataset...", total=None)
        counts = dataset.get_instance_counts()
        image_counts = dataset.get_image_counts()

    if not counts:
        console.print("[yellow]No annotations found in the dataset.[/yellow]")
        return

    # Collect all classes and splits
    all_classes: set[str] = set()
    all_splits: list[str] = []
    for split, class_counts in counts.items():
        all_splits.append(split)
        all_classes.update(class_counts.keys())

    # Sort splits in standard order
    split_order = {"train": 0, "val": 1, "test": 2}
    all_splits.sort(key=lambda s: split_order.get(s, 99))

    # Sort classes alphabetically
    sorted_classes = sorted(all_classes)

    # Create table
    title = f"Instance Statistics: {dataset_path.name} ({dataset.format.value})"
    table = Table(title=title)
    table.add_column("Class", style="cyan")
    for split in all_splits:
        table.add_column(split, justify="right", style="green")
    table.add_column("Total", justify="right", style="yellow bold")

    # Add rows for each class
    grand_totals = {split: 0 for split in all_splits}
    grand_total = 0

    for class_name in sorted_classes:
        row = [class_name]
        class_total = 0
        for split in all_splits:
            count = counts.get(split, {}).get(class_name, 0)
            row.append(str(count) if count > 0 else "-")
            class_total += count
            grand_totals[split] += count
        row.append(str(class_total))
        grand_total += class_total
        table.add_row(*row)

    # Add totals row
    table.add_section()
    totals_row = ["[bold]Total[/bold]"]
    for split in all_splits:
        totals_row.append(f"[bold]{grand_totals[split]}[/bold]")
    totals_row.append(f"[bold]{grand_total}[/bold]")
    table.add_row(*totals_row)

    console.print(table)

    # Build image stats line
    image_parts = []
    total_images = 0
    total_background = 0
    for split in all_splits:
        if split in image_counts:
            img_total = image_counts[split]["total"]
            img_bg = image_counts[split]["background"]
            total_images += img_total
            total_background += img_bg
            if img_bg > 0:
                image_parts.append(f"{split}: {img_total} ({img_bg} background)")
            else:
                image_parts.append(f"{split}: {img_total}")

    console.print(
        f"\n[green]Dataset: {dataset.format.value.upper()} | "
        f"Task: {dataset.task.value} | "
        f"Classes: {len(sorted_classes)} | "
        f"Total instances: {grand_total}[/green]"
    )

    if image_parts:
        console.print(f"[blue]Images: {' | '.join(image_parts)}[/blue]")


def _show_mask_stats(dataset: MaskDataset, dataset_path: Path) -> None:
    """Show statistics for mask datasets with pixel-level information.

    Args:
        dataset: The MaskDataset instance.
        dataset_path: Path to the dataset root.
    """
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
        transient=True,
    ) as progress:
        progress.add_task("Analyzing mask dataset...", total=None)
        pixel_counts = dataset.get_pixel_counts()
        image_presence = dataset.get_image_class_presence()
        image_counts = dataset.get_image_counts()

    # Get class mapping
    class_mapping = dataset.get_class_mapping()

    # Calculate total non-ignored pixels
    total_pixels = sum(
        count
        for class_id, count in pixel_counts.items()
        if class_id != dataset.ignore_index
    )
    ignored_pixels = pixel_counts.get(dataset.ignore_index, 0)

    # Calculate total images
    total_images = sum(ic["total"] for ic in image_counts.values())

    # Create table
    splits_str = ", ".join(dataset.splits) if dataset.splits else "unsplit"
    title = f"Class Statistics: {dataset_path.name} ({splits_str})"
    table = Table(title=title)
    table.add_column("Class", style="cyan")
    table.add_column("Total Pixels", justify="right", style="green")
    table.add_column("% Coverage", justify="right", style="magenta")
    table.add_column("Images With", justify="right", style="yellow")

    # Sort classes by class_id
    sorted_class_ids = sorted(class_mapping.keys())

    for class_id in sorted_class_ids:
        class_name = class_mapping[class_id]
        pixels = pixel_counts.get(class_id, 0)
        presence = image_presence.get(class_id, 0)

        # Calculate coverage percentage
        coverage = (pixels / total_pixels * 100) if total_pixels > 0 else 0.0

        table.add_row(
            class_name,
            f"{pixels:,}",
            f"{coverage:.1f}%",
            str(presence),
        )

    # Add ignored row if there are ignored pixels
    if ignored_pixels > 0:
        table.add_section()
        table.add_row(
            "[dim](ignored)[/dim]",
            f"[dim]{ignored_pixels:,}[/dim]",
            "[dim]-[/dim]",
            f"[dim]{total_images}[/dim]",
        )

    console.print(table)

    # Summary line
    console.print(f"\n[green]Dataset: {dataset_path}[/green]")
    console.print(
        f"[green]Format: {dataset.format.value.upper()} | "
        f"Task: {dataset.task.value}[/green]"
    )

    # Image counts per split
    image_parts = []
    for split in dataset.splits if dataset.splits else ["unsplit"]:
        if split in image_counts:
            image_parts.append(f"{split}: {image_counts[split]['total']}")

    if image_parts:
        console.print(f"[blue]Images: {' | '.join(image_parts)}[/blue]")


@app.command(name="view")
def view(
    dataset_path: Annotated[
        Path,
        typer.Option(
            "--dataset-path",
            "-d",
            help="Path to the dataset root directory.",
        ),
    ] = Path("."),
    split: Annotated[
        str | None,
        typer.Option(
            "--split",
            "-s",
            help="Specific split to view (train, val, test).",
        ),
    ] = None,
    max_classes: Annotated[
        int | None,
        typer.Option(
            "--max-classes",
            "-m",
            help="Maximum classes to show in grid (classification only).",
        ),
    ] = None,
    opacity: Annotated[
        float,
        typer.Option(
            "--opacity",
            "-o",
            help="Mask overlay opacity (0.0-1.0, mask datasets only).",
            min=0.0,
            max=1.0,
        ),
    ] = 0.5,
) -> None:
    """View annotated images in a dataset.

    Opens an interactive viewer to browse images with their annotations
    (bounding boxes and segmentation masks) overlaid.

    For classification datasets, shows a grid view with one image per class.

    Controls:
        - Right Arrow / N: Next image(s)
        - Left Arrow / P: Previous image(s)
        - Mouse Wheel: Zoom in/out (detection/segmentation only)
        - Mouse Drag: Pan when zoomed (detection/segmentation only)
        - R: Reset zoom / Reset to first images
        - T: Toggle annotations (detection/segmentation only)
        - Q / ESC: Quit viewer
    """
    # Resolve path and validate
    dataset_path = dataset_path.resolve()
    if not dataset_path.exists():
        console.print(f"[red]Error: Path does not exist: {dataset_path}[/red]")
        raise typer.Exit(1)
    if not dataset_path.is_dir():
        console.print(f"[red]Error: Path is not a directory: {dataset_path}[/red]")
        raise typer.Exit(1)

    # Detect dataset
    dataset = _detect_dataset(dataset_path)
    if not dataset:
        console.print(
            f"[red]Error: No YOLO or COCO dataset found at {dataset_path}[/red]\n"
            "[yellow]Ensure the path points to a dataset root containing "
            "data.yaml (YOLO) or annotations/ folder (COCO).[/yellow]"
        )
        raise typer.Exit(1)

    # Validate split if specified
    if split and split not in dataset.splits:
        available = ", ".join(dataset.splits) if dataset.splits else "none"
        console.print(
            f"[red]Error: Split '{split}' not found in dataset.[/red]\n"
            f"[yellow]Available splits: {available}[/yellow]"
        )
        raise typer.Exit(1)

    # Generate consistent colors for each class
    class_colors = _generate_class_colors(dataset.class_names)

    # Handle mask datasets with overlay viewer
    if dataset.format == DatasetFormat.MASK:
        assert isinstance(dataset, MaskDataset)
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Loading images...", total=None)
            image_paths = dataset.get_image_paths(split)

        if not image_paths:
            console.print("[yellow]No images found in the dataset.[/yellow]")
            return

        console.print(
            f"[green]Found {len(image_paths)} images. "
            f"Opening mask viewer...[/green]\n"
            "[dim]Controls: \u2190 / \u2192 or P / N to navigate, "
            "Mouse wheel to zoom, Drag to pan, R to reset, T to toggle overlay, "
            "Q / ESC to quit[/dim]"
        )

        viewer = _MaskViewer(
            image_paths=image_paths,
            dataset=dataset,
            class_colors=class_colors,
            window_name=f"Argus Mask Viewer - {dataset_path.name}",
            opacity=opacity,
        )
        viewer.run()
        console.print("[green]Viewer closed.[/green]")
        return

    # Handle classification datasets with grid viewer
    if dataset.task == TaskType.CLASSIFICATION:
        # Use first split if specified, otherwise let get_images_by_class handle it
        view_split = split if split else (dataset.splits[0] if dataset.splits else None)

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Loading images by class...", total=None)
            images_by_class = dataset.get_images_by_class(view_split)

        total_images = sum(len(imgs) for imgs in images_by_class.values())
        if total_images == 0:
            console.print("[yellow]No images found in the dataset.[/yellow]")
            return

        num_classes = len(dataset.class_names)
        display_classes = min(num_classes, max_classes) if max_classes else num_classes

        console.print(
            f"[green]Found {total_images} images across {num_classes} classes "
            f"(showing {display_classes}). Opening grid viewer...[/green]\n"
            "[dim]Controls: ← / → or P / N to navigate all classes, "
            "R to reset, Q / ESC to quit[/dim]"
        )

        viewer = _ClassificationGridViewer(
            images_by_class=images_by_class,
            class_names=dataset.class_names,
            class_colors=class_colors,
            window_name=f"Argus Classification Viewer - {dataset_path.name}",
            max_classes=max_classes,
        )
        viewer.run()
    else:
        # Detection/Segmentation viewer
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Loading images...", total=None)
            image_paths = dataset.get_image_paths(split)

        if not image_paths:
            console.print("[yellow]No images found in the dataset.[/yellow]")
            return

        console.print(
            f"[green]Found {len(image_paths)} images. "
            f"Opening viewer...[/green]\n"
            "[dim]Controls: ← / → or P / N to navigate, "
            "Mouse wheel to zoom, Drag to pan, R to reset, T to toggle annotations, "
            "Q / ESC to quit[/dim]"
        )

        viewer = _ImageViewer(
            image_paths=image_paths,
            dataset=dataset,
            class_colors=class_colors,
            window_name=f"Argus Viewer - {dataset_path.name}",
        )
        viewer.run()

    console.print("[green]Viewer closed.[/green]")


@app.command(name="split")
def split_dataset(
    dataset_path: Annotated[
        Path,
        typer.Option(
            "--dataset-path",
            "-d",
            help="Path to the dataset root directory.",
        ),
    ] = Path("."),
    output_path: Annotated[
        Path,
        typer.Option(
            "--output-path",
            "-o",
            help="Directory to write the split dataset.",
        ),
    ] = Path("splits"),
    ratio: Annotated[
        str,
        typer.Option(
            "--ratio",
            "-r",
            help="Train/val/test ratio (e.g. 0.8,0.1,0.1).",
        ),
    ] = "0.8,0.1,0.1",
    seed: Annotated[
        int,
        typer.Option(
            "--seed",
            help="Random seed for deterministic splitting.",
        ),
    ] = 42,
) -> None:
    """Split an unsplit dataset into train/val/test."""
    dataset_path = dataset_path.resolve()
    if not dataset_path.exists():
        console.print(f"[red]Error: Path does not exist: {dataset_path}[/red]")
        raise typer.Exit(1)
    if not dataset_path.is_dir():
        console.print(f"[red]Error: Path is not a directory: {dataset_path}[/red]")
        raise typer.Exit(1)

    dataset = _detect_dataset(dataset_path)
    if not dataset:
        console.print(
            f"[red]Error: No YOLO or COCO dataset found at {dataset_path}[/red]\n"
            "[yellow]Ensure the path points to a dataset root containing "
            "data.yaml (YOLO) or annotations/ folder (COCO).[/yellow]"
        )
        raise typer.Exit(1)

    try:
        ratios = parse_ratio(ratio)
    except ValueError as exc:
        console.print(f"[red]Error: {exc}[/red]")
        raise typer.Exit(1) from exc

    if not output_path.is_absolute():
        output_path = dataset_path / output_path
    output_path = output_path.resolve()

    if isinstance(dataset, YOLODataset):
        if dataset.splits:
            console.print(
                "[red]Error: Dataset already has splits. "
                "Use an unsplit dataset to run split.[/red]"
            )
            raise typer.Exit(1)

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Creating YOLO splits...", total=None)
            try:
                counts = split_yolo_dataset(dataset, output_path, ratios, True, seed)
            except ValueError as exc:
                console.print(f"[red]Error: {exc}[/red]")
                raise typer.Exit(1) from exc
    else:
        coco_dataset = dataset
        if not isinstance(coco_dataset, COCODataset):
            console.print("[red]Error: Unsupported dataset type.[/red]")
            raise typer.Exit(1)
        if not is_coco_unsplit(coco_dataset.annotation_files):
            console.print(
                "[red]Error: Dataset already has splits. "
                "Use an unsplit dataset to run split.[/red]"
            )
            raise typer.Exit(1)
        if not coco_dataset.annotation_files:
            console.print("[red]Error: No annotation files found.[/red]")
            raise typer.Exit(1)
        annotation_file = coco_dataset.annotation_files[0]

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
            transient=True,
        ) as progress:
            progress.add_task("Creating COCO splits...", total=None)
            try:
                counts = split_coco_dataset(
                    coco_dataset,
                    annotation_file,
                    output_path,
                    ratios,
                    True,
                    seed,
                )
            except ValueError as exc:
                console.print(f"[red]Error: {exc}[/red]")
                raise typer.Exit(1) from exc

    console.print(
        "[green]Split complete.[/green] "
        f"Train: {counts['train']}, Val: {counts['val']}, Test: {counts['test']}."
    )


@app.command(name="convert")
def convert_dataset(
    input_path: Annotated[
        Path,
        typer.Option(
            "--input-path",
            "-i",
            help="Path to the source dataset.",
        ),
    ] = Path("."),
    output_path: Annotated[
        Path,
        typer.Option(
            "--output-path",
            "-o",
            help="Output directory for converted dataset.",
        ),
    ] = Path("converted"),
    to_format: Annotated[
        str,
        typer.Option(
            "--to",
            help="Target format (currently only 'yolo-seg' is supported).",
        ),
    ] = "yolo-seg",
    epsilon_factor: Annotated[
        float,
        typer.Option(
            "--epsilon-factor",
            "-e",
            help="Polygon simplification factor (Douglas-Peucker algorithm).",
            min=0.0,
            max=1.0,
        ),
    ] = 0.005,
    min_area: Annotated[
        float,
        typer.Option(
            "--min-area",
            "-a",
            help="Minimum contour area in pixels to include.",
            min=0.0,
        ),
    ] = 100.0,
) -> None:
    """Convert a dataset from one format to another.

    Currently supports converting MaskDataset to YOLO segmentation format.

    Example:
        uvx argus-cv convert -i /path/to/masks -o /path/to/output --to yolo-seg
    """
    # Validate format
    if to_format != "yolo-seg":
        console.print(
            f"[red]Error: Unsupported target format '{to_format}'.[/red]\n"
            "[yellow]Currently only 'yolo-seg' is supported.[/yellow]"
        )
        raise typer.Exit(1)

    # Resolve and validate input path
    input_path = input_path.resolve()
    if not input_path.exists():
        console.print(f"[red]Error: Path does not exist: {input_path}[/red]")
        raise typer.Exit(1)
    if not input_path.is_dir():
        console.print(f"[red]Error: Path is not a directory: {input_path}[/red]")
        raise typer.Exit(1)

    # Detect source dataset - must be MaskDataset for yolo-seg conversion
    dataset = MaskDataset.detect(input_path)
    if not dataset:
        console.print(
            f"[red]Error: No MaskDataset found at {input_path}[/red]\n"
            "[yellow]Ensure the path contains images/ + masks/ directories "
            "(or equivalent patterns like img/+gt/ or leftImg8bit/+gtFine/).[/yellow]"
        )
        raise typer.Exit(1)

    # Resolve output path
    if not output_path.is_absolute():
        output_path = input_path.parent / output_path
    output_path = output_path.resolve()

    # Check if output already exists
    if output_path.exists() and any(output_path.iterdir()):
        console.print(
            f"[red]Error: Output directory already exists and is not empty: "
            f"{output_path}[/red]"
        )
        raise typer.Exit(1)

    # Show conversion info
    console.print("[cyan]Converting MaskDataset to YOLO segmentation format[/cyan]")
    console.print(f"  Source: {input_path}")
    console.print(f"  Output: {output_path}")
    console.print(f"  Classes: {dataset.num_classes}")
    splits_str = ", ".join(dataset.splits) if dataset.splits else "unsplit"
    console.print(f"  Splits: {splits_str}")
    console.print()

    # Run conversion with progress bar
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Processing images...", total=None)

        def update_progress(current: int, total: int) -> None:
            progress.update(task, completed=current, total=total)

        try:
            stats = convert_mask_to_yolo_seg(
                dataset=dataset,
                output_path=output_path,
                epsilon_factor=epsilon_factor,
                min_area=min_area,
                progress_callback=update_progress,
            )
        except Exception as exc:
            console.print(f"[red]Error during conversion: {exc}[/red]")
            raise typer.Exit(1) from exc

    # Show results
    console.print()
    console.print("[green]Conversion complete![/green]")
    console.print(f"  Images processed: {stats['images']}")
    console.print(f"  Labels created: {stats['labels']}")
    console.print(f"  Polygons extracted: {stats['polygons']}")

    if stats["skipped"] > 0:
        skipped = stats["skipped"]
        console.print(f"  [yellow]Skipped: {skipped} (no mask or empty)[/yellow]")
    if stats["warnings"] > 0:
        console.print(f"  [yellow]Warnings: {stats['warnings']}[/yellow]")

    console.print(f"\n[cyan]Output dataset: {output_path}[/cyan]")


class _ImageViewer:
    """Interactive image viewer with zoom and pan support."""

    def __init__(
        self,
        image_paths: list[Path],
        dataset: Dataset,
        class_colors: dict[str, tuple[int, int, int]],
        window_name: str,
    ):
        self.image_paths = image_paths
        self.dataset = dataset
        self.class_colors = class_colors
        self.window_name = window_name

        self.current_idx = 0
        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0

        # Mouse state for panning
        self.dragging = False
        self.drag_start_x = 0
        self.drag_start_y = 0
        self.pan_start_x = 0.0
        self.pan_start_y = 0.0

        # Current image cache
        self.current_img: np.ndarray | None = None
        self.annotated_img: np.ndarray | None = None

        # Annotation visibility toggle
        self.show_annotations = True

    def _load_current_image(self) -> bool:
        """Load and annotate the current image."""
        image_path = self.image_paths[self.current_idx]
        annotations = self.dataset.get_annotations_for_image(image_path)

        img = cv2.imread(str(image_path))
        if img is None:
            return False

        self.current_img = img
        self.annotated_img = _draw_annotations(
            img.copy(), annotations, self.class_colors
        )
        return True

    def _get_display_image(self) -> np.ndarray:
        """Get the image transformed for current zoom/pan."""
        if self.annotated_img is None:
            return np.zeros((480, 640, 3), dtype=np.uint8)

        if self.show_annotations:
            img = self.annotated_img
        elif self.current_img is not None:
            img = self.current_img
        else:
            img = self.annotated_img
        h, w = img.shape[:2]

        if self.zoom == 1.0 and self.pan_x == 0.0 and self.pan_y == 0.0:
            display = img.copy()
        else:
            # Calculate the visible region
            view_w = int(w / self.zoom)
            view_h = int(h / self.zoom)

            # Center point with pan offset
            cx = w / 2 + self.pan_x
            cy = h / 2 + self.pan_y

            # Calculate crop bounds
            x1 = int(max(0, cx - view_w / 2))
            y1 = int(max(0, cy - view_h / 2))
            x2 = int(min(w, x1 + view_w))
            y2 = int(min(h, y1 + view_h))

            # Adjust if we hit boundaries
            if x2 - x1 < view_w:
                x1 = max(0, x2 - view_w)
            if y2 - y1 < view_h:
                y1 = max(0, y2 - view_h)

            # Crop and resize
            cropped = img[y1:y2, x1:x2]
            display = cv2.resize(cropped, (w, h), interpolation=cv2.INTER_LINEAR)

        # Add info overlay
        image_path = self.image_paths[self.current_idx]
        idx = self.current_idx + 1
        total = len(self.image_paths)
        info_text = f"[{idx}/{total}] {image_path.name}"
        if self.zoom > 1.0:
            info_text += f" (Zoom: {self.zoom:.1f}x)"
        if not self.show_annotations:
            info_text += " [Annotations: OFF]"

        cv2.putText(
            display,
            info_text,
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            2,
        )
        cv2.putText(
            display, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1
        )

        return display

    def _mouse_callback(
        self, event: int, x: int, y: int, flags: int, param: None
    ) -> None:
        """Handle mouse events for zoom and pan."""
        if event == cv2.EVENT_MOUSEWHEEL:
            # Zoom in/out
            if flags > 0:
                self.zoom = min(10.0, self.zoom * 1.2)
            else:
                self.zoom = max(1.0, self.zoom / 1.2)

            # Reset pan if zoomed out to 1x
            if self.zoom == 1.0:
                self.pan_x = 0.0
                self.pan_y = 0.0

        elif event == cv2.EVENT_LBUTTONDOWN:
            self.dragging = True
            self.drag_start_x = x
            self.drag_start_y = y
            self.pan_start_x = self.pan_x
            self.pan_start_y = self.pan_y

        elif event == cv2.EVENT_MOUSEMOVE and self.dragging:
            if self.zoom > 1.0 and self.annotated_img is not None:
                h, w = self.annotated_img.shape[:2]
                # Calculate pan delta (inverted for natural feel)
                dx = (self.drag_start_x - x) / self.zoom
                dy = (self.drag_start_y - y) / self.zoom

                # Update pan with limits
                max_pan_x = w * (1 - 1 / self.zoom) / 2
                max_pan_y = h * (1 - 1 / self.zoom) / 2

                self.pan_x = max(-max_pan_x, min(max_pan_x, self.pan_start_x + dx))
                self.pan_y = max(-max_pan_y, min(max_pan_y, self.pan_start_y + dy))

        elif event == cv2.EVENT_LBUTTONUP:
            self.dragging = False

    def _reset_view(self) -> None:
        """Reset zoom and pan to default."""
        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0

    def _next_image(self) -> None:
        """Go to next image."""
        self.current_idx = (self.current_idx + 1) % len(self.image_paths)
        self._reset_view()

    def _prev_image(self) -> None:
        """Go to previous image."""
        self.current_idx = (self.current_idx - 1) % len(self.image_paths)
        self._reset_view()

    def run(self) -> None:
        """Run the interactive viewer."""
        cv2.namedWindow(self.window_name, cv2.WINDOW_AUTOSIZE)
        cv2.setMouseCallback(self.window_name, self._mouse_callback)

        while True:
            # Load image if needed
            if self.annotated_img is None and not self._load_current_image():
                console.print(
                    f"[yellow]Warning: Could not load "
                    f"{self.image_paths[self.current_idx]}[/yellow]"
                )
                self._next_image()
                continue

            # Display image
            display = self._get_display_image()
            cv2.imshow(self.window_name, display)

            # Wait for input (short timeout for smooth panning)
            key = cv2.waitKey(30) & 0xFF

            # Handle keyboard input
            if key == ord("q") or key == 27:  # Q or ESC
                break
            elif key == ord("n") or key == 83 or key == 3:  # N or Right arrow
                self.annotated_img = None
                self._next_image()
            elif key == ord("p") or key == 81 or key == 2:  # P or Left arrow
                self.annotated_img = None
                self._prev_image()
            elif key == ord("r"):  # R to reset zoom
                self._reset_view()
            elif key == ord("t"):  # T to toggle annotations
                self.show_annotations = not self.show_annotations

        cv2.destroyAllWindows()


class _ClassificationGridViewer:
    """Grid viewer for classification datasets showing one image per class."""

    def __init__(
        self,
        images_by_class: dict[str, list[Path]],
        class_names: list[str],
        class_colors: dict[str, tuple[int, int, int]],
        window_name: str,
        max_classes: int | None = None,
        tile_size: int = 300,
    ):
        # Limit classes if max_classes specified
        if max_classes and len(class_names) > max_classes:
            self.class_names = class_names[:max_classes]
        else:
            self.class_names = class_names

        self.images_by_class = {
            cls: images_by_class.get(cls, []) for cls in self.class_names
        }
        self.class_colors = class_colors
        self.window_name = window_name
        self.tile_size = tile_size

        # Global image index (same for all classes)
        self.current_index = 0

        # Calculate max images across all classes
        self.max_images = (
            max(len(imgs) for imgs in self.images_by_class.values())
            if self.images_by_class
            else 0
        )

        # Calculate grid layout
        self.cols, self.rows = self._calculate_grid_layout()

    def _calculate_grid_layout(self) -> tuple[int, int]:
        """Calculate optimal grid layout based on number of classes."""
        n = len(self.class_names)
        if n <= 0:
            return 1, 1

        # Try to make a roughly square grid
        import math

        cols = int(math.ceil(math.sqrt(n)))
        rows = int(math.ceil(n / cols))
        return cols, rows

    def _create_tile(
        self, class_name: str, image_path: Path | None, index: int, total: int
    ) -> np.ndarray:
        """Create a single tile for a class."""
        tile = np.zeros((self.tile_size, self.tile_size, 3), dtype=np.uint8)

        if image_path is not None and image_path.exists():
            # Load and resize image
            img = cv2.imread(str(image_path))
            if img is not None:
                # Resize maintaining aspect ratio
                h, w = img.shape[:2]
                scale = min(self.tile_size / w, self.tile_size / h)
                new_w = int(w * scale)
                new_h = int(h * scale)
                resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

                # Center in tile
                x_offset = (self.tile_size - new_w) // 2
                y_offset = (self.tile_size - new_h) // 2
                tile[y_offset : y_offset + new_h, x_offset : x_offset + new_w] = resized

        # Draw label at top: "class_name (N/M)"
        if image_path is not None:
            label = f"{class_name} ({index + 1}/{total})"
        else:
            label = f"{class_name} (-/{total})"

        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        thickness = 1
        (label_w, label_h), baseline = cv2.getTextSize(
            label, font, font_scale, thickness
        )

        # Semi-transparent background for label
        overlay = tile.copy()
        label_bg_height = label_h + baseline + 10
        cv2.rectangle(overlay, (0, 0), (self.tile_size, label_bg_height), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.6, tile, 0.4, 0, tile)

        cv2.putText(
            tile,
            label,
            (5, label_h + 5),
            font,
            font_scale,
            (255, 255, 255),
            thickness,
        )

        # Draw thin border
        border_end = self.tile_size - 1
        cv2.rectangle(tile, (0, 0), (border_end, border_end), (80, 80, 80), 1)

        return tile

    def _compose_grid(self) -> np.ndarray:
        """Compose all tiles into a single grid image."""
        grid_h = self.rows * self.tile_size
        grid_w = self.cols * self.tile_size
        grid = np.zeros((grid_h, grid_w, 3), dtype=np.uint8)

        for i, class_name in enumerate(self.class_names):
            row = i // self.cols
            col = i % self.cols

            images = self.images_by_class[class_name]
            total = len(images)

            # Use global index - show black tile if class doesn't have this image
            if self.current_index < total:
                image_path = images[self.current_index]
                display_index = self.current_index
            else:
                image_path = None
                display_index = self.current_index

            tile = self._create_tile(class_name, image_path, display_index, total)

            y_start = row * self.tile_size
            x_start = col * self.tile_size
            y_end = y_start + self.tile_size
            x_end = x_start + self.tile_size
            grid[y_start:y_end, x_start:x_end] = tile

        return grid

    def _next_images(self) -> None:
        """Advance to next image index."""
        if self.max_images > 0:
            self.current_index = min(self.current_index + 1, self.max_images - 1)

    def _prev_images(self) -> None:
        """Go back to previous image index."""
        self.current_index = max(self.current_index - 1, 0)

    def _reset_indices(self) -> None:
        """Reset to first image."""
        self.current_index = 0

    def run(self) -> None:
        """Run the interactive grid viewer."""
        if not self.class_names:
            console.print("[yellow]No classes to display.[/yellow]")
            return

        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)

        while True:
            # Compose and display grid
            grid = self._compose_grid()
            cv2.imshow(self.window_name, grid)

            # Wait for input
            key = cv2.waitKey(30) & 0xFF

            # Handle keyboard input
            if key == ord("q") or key == 27:  # Q or ESC
                break
            elif key == ord("n") or key == 83 or key == 3:  # N or Right arrow
                self._next_images()
            elif key == ord("p") or key == 81 or key == 2:  # P or Left arrow
                self._prev_images()
            elif key == ord("r"):  # R to reset
                self._reset_indices()

        cv2.destroyAllWindows()


class _MaskViewer:
    """Interactive viewer for semantic mask datasets with colored overlay."""

    def __init__(
        self,
        image_paths: list[Path],
        dataset: MaskDataset,
        class_colors: dict[str, tuple[int, int, int]],
        window_name: str,
        opacity: float = 0.5,
    ):
        self.image_paths = image_paths
        self.dataset = dataset
        self.class_colors = class_colors
        self.window_name = window_name
        self.opacity = opacity

        self.current_idx = 0
        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0

        # Mouse state for panning
        self.dragging = False
        self.drag_start_x = 0
        self.drag_start_y = 0
        self.pan_start_x = 0.0
        self.pan_start_y = 0.0

        # Current image cache
        self.current_img: np.ndarray | None = None
        self.overlay_img: np.ndarray | None = None

        # Overlay visibility toggle
        self.show_overlay = True

        # Build class_id to color mapping
        self._id_to_color: dict[int, tuple[int, int, int]] = {}
        class_mapping = dataset.get_class_mapping()
        for class_id, class_name in class_mapping.items():
            if class_name in class_colors:
                self._id_to_color[class_id] = class_colors[class_name]

    def _load_current_image(self) -> bool:
        """Load current image and create mask overlay."""
        image_path = self.image_paths[self.current_idx]

        img = cv2.imread(str(image_path))
        if img is None:
            return False

        mask = self.dataset.load_mask(image_path)
        if mask is None:
            console.print(f"[yellow]Warning: No mask for {image_path}[/yellow]")
            self.current_img = img
            self.overlay_img = img.copy()
            return True

        # Validate dimensions
        if img.shape[:2] != mask.shape[:2]:
            console.print(
                f"[red]Error: Dimension mismatch for {image_path.name}: "
                f"image={img.shape[:2]}, mask={mask.shape[:2]}[/red]"
            )
            return False

        self.current_img = img
        self.overlay_img = self._create_overlay(img, mask)
        return True

    def _create_overlay(self, img: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Create colored overlay from mask.

        Args:
            img: Original image (BGR).
            mask: Grayscale mask with class IDs.

        Returns:
            Image with colored mask overlay.
        """
        # Create colored mask
        h, w = mask.shape
        colored_mask = np.zeros((h, w, 3), dtype=np.uint8)

        for class_id, color in self._id_to_color.items():
            colored_mask[mask == class_id] = color

        # Blend with original image
        # Ignore pixels are fully transparent (not blended)
        ignore_mask = mask == self.dataset.ignore_index
        alpha = np.ones((h, w, 1), dtype=np.float32) * self.opacity
        alpha[ignore_mask] = 0.0

        # Blend: result = img * (1 - alpha) + colored_mask * alpha
        blended = (
            img.astype(np.float32) * (1 - alpha)
            + colored_mask.astype(np.float32) * alpha
        )
        return blended.astype(np.uint8)

    def _get_display_image(self) -> np.ndarray:
        """Get the image transformed for current zoom/pan."""
        if self.overlay_img is None:
            return np.zeros((480, 640, 3), dtype=np.uint8)

        if self.show_overlay:
            img = self.overlay_img
        elif self.current_img is not None:
            img = self.current_img
        else:
            img = self.overlay_img

        h, w = img.shape[:2]

        if self.zoom == 1.0 and self.pan_x == 0.0 and self.pan_y == 0.0:
            display = img.copy()
        else:
            # Calculate the visible region
            view_w = int(w / self.zoom)
            view_h = int(h / self.zoom)

            # Center point with pan offset
            cx = w / 2 + self.pan_x
            cy = h / 2 + self.pan_y

            # Calculate crop bounds
            x1 = int(max(0, cx - view_w / 2))
            y1 = int(max(0, cy - view_h / 2))
            x2 = int(min(w, x1 + view_w))
            y2 = int(min(h, y1 + view_h))

            # Adjust if we hit boundaries
            if x2 - x1 < view_w:
                x1 = max(0, x2 - view_w)
            if y2 - y1 < view_h:
                y1 = max(0, y2 - view_h)

            # Crop and resize
            cropped = img[y1:y2, x1:x2]
            display = cv2.resize(cropped, (w, h), interpolation=cv2.INTER_LINEAR)

        # Add info overlay
        image_path = self.image_paths[self.current_idx]
        idx = self.current_idx + 1
        total = len(self.image_paths)
        info_text = f"[{idx}/{total}] {image_path.name}"
        if self.zoom > 1.0:
            info_text += f" (Zoom: {self.zoom:.1f}x)"
        if not self.show_overlay:
            info_text += " [Overlay: OFF]"

        cv2.putText(
            display,
            info_text,
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            2,
        )
        cv2.putText(
            display, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1
        )

        return display

    def _mouse_callback(
        self, event: int, x: int, y: int, flags: int, param: None
    ) -> None:
        """Handle mouse events for zoom and pan."""
        if event == cv2.EVENT_MOUSEWHEEL:
            # Zoom in/out
            if flags > 0:
                self.zoom = min(10.0, self.zoom * 1.2)
            else:
                self.zoom = max(1.0, self.zoom / 1.2)

            # Reset pan if zoomed out to 1x
            if self.zoom == 1.0:
                self.pan_x = 0.0
                self.pan_y = 0.0

        elif event == cv2.EVENT_LBUTTONDOWN:
            self.dragging = True
            self.drag_start_x = x
            self.drag_start_y = y
            self.pan_start_x = self.pan_x
            self.pan_start_y = self.pan_y

        elif event == cv2.EVENT_MOUSEMOVE and self.dragging:
            if self.zoom > 1.0 and self.overlay_img is not None:
                h, w = self.overlay_img.shape[:2]
                # Calculate pan delta (inverted for natural feel)
                dx = (self.drag_start_x - x) / self.zoom
                dy = (self.drag_start_y - y) / self.zoom

                # Update pan with limits
                max_pan_x = w * (1 - 1 / self.zoom) / 2
                max_pan_y = h * (1 - 1 / self.zoom) / 2

                self.pan_x = max(-max_pan_x, min(max_pan_x, self.pan_start_x + dx))
                self.pan_y = max(-max_pan_y, min(max_pan_y, self.pan_start_y + dy))

        elif event == cv2.EVENT_LBUTTONUP:
            self.dragging = False

    def _reset_view(self) -> None:
        """Reset zoom and pan to default."""
        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0

    def _next_image(self) -> None:
        """Go to next image."""
        self.current_idx = (self.current_idx + 1) % len(self.image_paths)
        self._reset_view()

    def _prev_image(self) -> None:
        """Go to previous image."""
        self.current_idx = (self.current_idx - 1) % len(self.image_paths)
        self._reset_view()

    def run(self) -> None:
        """Run the interactive viewer."""
        cv2.namedWindow(self.window_name, cv2.WINDOW_AUTOSIZE)
        cv2.setMouseCallback(self.window_name, self._mouse_callback)

        while True:
            # Load image if needed
            if self.overlay_img is None and not self._load_current_image():
                console.print(
                    f"[yellow]Warning: Could not load "
                    f"{self.image_paths[self.current_idx]}[/yellow]"
                )
                self._next_image()
                continue

            # Display image
            display = self._get_display_image()
            cv2.imshow(self.window_name, display)

            # Wait for input (short timeout for smooth panning)
            key = cv2.waitKey(30) & 0xFF

            # Handle keyboard input
            if key == ord("q") or key == 27:  # Q or ESC
                break
            elif key == ord("n") or key == 83 or key == 3:  # N or Right arrow
                self.overlay_img = None
                self._next_image()
            elif key == ord("p") or key == 81 or key == 2:  # P or Left arrow
                self.overlay_img = None
                self._prev_image()
            elif key == ord("r"):  # R to reset zoom
                self._reset_view()
            elif key == ord("t"):  # T to toggle overlay
                self.show_overlay = not self.show_overlay

        cv2.destroyAllWindows()


def _generate_class_colors(class_names: list[str]) -> dict[str, tuple[int, int, int]]:
    """Generate consistent colors for each class name.

    Args:
        class_names: List of class names.

    Returns:
        Dictionary mapping class name to BGR color tuple.
    """
    colors: dict[str, tuple[int, int, int]] = {}

    for name in class_names:
        # Generate a consistent hash-based color
        hash_val = int(hashlib.md5(name.encode()).hexdigest()[:6], 16)
        r = (hash_val >> 16) & 0xFF
        g = (hash_val >> 8) & 0xFF
        b = hash_val & 0xFF

        # Ensure colors are bright enough to be visible
        min_brightness = 100
        r = max(r, min_brightness)
        g = max(g, min_brightness)
        b = max(b, min_brightness)

        colors[name] = (b, g, r)  # BGR for OpenCV

    return colors


def _draw_annotations(
    img: np.ndarray,
    annotations: list[dict],
    class_colors: dict[str, tuple[int, int, int]],
) -> np.ndarray:
    """Draw annotations on an image.

    Args:
        img: OpenCV image (BGR).
        annotations: List of annotation dicts.
        class_colors: Dictionary mapping class name to BGR color.

    Returns:
        Image with annotations drawn.
    """
    default_color = (0, 255, 0)  # Green default

    for ann in annotations:
        class_name = ann["class_name"]
        color = class_colors.get(class_name, default_color)
        bbox = ann.get("bbox")
        polygon = ann.get("polygon")

        # Draw polygon if available (segmentation)
        if polygon:
            pts = np.array(polygon, dtype=np.int32)
            cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)
            # Draw semi-transparent fill
            overlay = img.copy()
            cv2.fillPoly(overlay, [pts], color)
            cv2.addWeighted(overlay, 0.3, img, 0.7, 0, img)
            # Draw small points at polygon vertices
            for pt in pts:
                cv2.circle(img, tuple(pt), radius=3, color=color, thickness=-1)

        # Draw bounding box (only for detection, not segmentation)
        if bbox and not polygon:
            x, y, w, h = bbox
            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)

            # Draw label background
            label = class_name
            (label_w, label_h), baseline = cv2.getTextSize(
                label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
            )
            cv2.rectangle(
                img,
                (x1, y1 - label_h - baseline - 5),
                (x1 + label_w + 5, y1),
                color,
                -1,
            )
            # Draw label text
            cv2.putText(
                img,
                label,
                (x1 + 2, y1 - baseline - 2),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                1,
            )

    return img


def _discover_datasets(root_path: Path, max_depth: int) -> list[Dataset]:
    """Discover all datasets under the root path.

    Args:
        root_path: Root directory to search.
        max_depth: Maximum depth to traverse.

    Returns:
        List of discovered Dataset instances.
    """
    datasets: list[Dataset] = []
    visited_paths: set[Path] = set()

    def _walk_directory(current_path: Path, depth: int) -> None:
        """Recursively walk directories and detect datasets."""
        if depth > max_depth:
            return

        if not current_path.is_dir():
            return

        # Normalize path to avoid duplicates
        resolved_path = current_path.resolve()
        if resolved_path in visited_paths:
            return
        visited_paths.add(resolved_path)

        # Try to detect datasets at this level
        dataset = _detect_dataset(current_path)
        if dataset:
            # Check if we already have a dataset for this path
            if not any(d.path.resolve() == resolved_path for d in datasets):
                datasets.append(dataset)
            # Don't recurse into detected datasets to avoid duplicates
            return

        # Recurse into subdirectories
        try:
            for entry in current_path.iterdir():
                if entry.is_dir() and not entry.name.startswith("."):
                    _walk_directory(entry, depth + 1)
        except PermissionError:
            pass  # Skip directories we can't access

    _walk_directory(root_path, 0)

    # Sort datasets by path for consistent output
    datasets.sort(key=lambda d: str(d.path))

    return datasets


def _detect_dataset(path: Path) -> Dataset | None:
    """Try to detect a dataset at the given path.

    Detection priority: YOLO -> COCO -> MaskDataset
    """
    # Try YOLO first (more specific patterns - requires data.yaml)
    dataset = YOLODataset.detect(path)
    if dataset:
        return dataset

    # Try COCO (requires annotations/*.json)
    dataset = COCODataset.detect(path)
    if dataset:
        return dataset

    # Try MaskDataset (directory structure based)
    dataset = MaskDataset.detect(path)
    if dataset:
        return dataset

    return None


if __name__ == "__main__":
    app()
