"""
Base classes for augmentation transforms.

Provides abstract base class and result dataclass for all augmentation
implementations.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any

import numpy as np


@dataclass
class AugmentationResult:
    """
    Result of applying an augmentation.

    Attributes:
        image: The augmented image as numpy array (H, W, C).
        bboxes: Updated bounding boxes if geometric transform was applied.
                Format: (N, 5) array with [class_id, x_center, y_center, width, height].
        transform_matrix: The transformation matrix if applicable (for bbox adjustment).
        applied: Whether the augmentation was actually applied.
        metadata: Additional metadata about the augmentation.
    """

    image: np.ndarray
    bboxes: np.ndarray | None = None
    transform_matrix: np.ndarray | None = None
    applied: bool = True
    metadata: dict[str, Any] | None = None


class BaseAugmentation(ABC):
    """
    Abstract base class for all augmentations.

    Subclasses must implement:
        - _validate_params(): Validate augmentation parameters
        - apply(): Apply the augmentation to an image

    Class attributes:
        name: Human-readable name of the augmentation.
        affects_geometry: True if this augmentation modifies bbox coordinates.
    """

    name: str = "base"
    affects_geometry: bool = False

    def __init__(self, params: dict[str, Any]) -> None:
        """
        Initialize augmentation with parameters.

        Args:
            params: Dictionary of augmentation-specific parameters.
        """
        self.params = params
        self._validate_params()

    @abstractmethod
    def _validate_params(self) -> None:
        """
        Validate augmentation parameters.

        Raises:
            ValueError: If parameters are invalid.
        """
        pass

    @abstractmethod
    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        """
        Apply augmentation to image.

        IMPORTANT: Implementations must NOT modify the input image or bboxes.
        Always create copies before modifying.

        Args:
            image: Input image as numpy array (H, W, C) with dtype uint8.
            bboxes: Optional bounding boxes in YOLO format (N, 5) array.
                    Each row: [class_id, x_center, y_center, width, height].
                    Coordinates are normalized to 0-1 range.
            rng: Random number generator for reproducibility.
                 If None, a new generator should be created.

        Returns:
            AugmentationResult with augmented image and optionally updated bboxes.
        """
        pass

    def get_preview_params(self) -> dict[str, Any]:
        """
        Get parameters optimized for preview display.

        Override this method to provide parameters that produce
        clearly visible effects for preview/demo purposes.

        Returns:
            Dictionary of preview parameters.
        """
        return dict(self.params)
