"""
Augmentation pipeline module.

Orchestrates multiple augmentations with proper ordering and
provides preview functionality.
"""

from typing import Any

import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation
from shared.augmentation.config import AugmentationConfig, AugmentationParams
from shared.augmentation.transforms.blur import GaussianBlur, MotionBlur
from shared.augmentation.transforms.degradation import EdgeDamage, Stain, Wrinkle
from shared.augmentation.transforms.geometric import PerspectiveWarp
from shared.augmentation.transforms.lighting import LightingVariation, Shadow
from shared.augmentation.transforms.noise import GaussianNoise, SaltPepper
from shared.augmentation.transforms.texture import PaperTexture, ScannerArtifacts

# Registry of augmentation classes
AUGMENTATION_REGISTRY: dict[str, type[BaseAugmentation]] = {
    "perspective_warp": PerspectiveWarp,
    "wrinkle": Wrinkle,
    "edge_damage": EdgeDamage,
    "stain": Stain,
    "lighting_variation": LightingVariation,
    "shadow": Shadow,
    "gaussian_blur": GaussianBlur,
    "motion_blur": MotionBlur,
    "gaussian_noise": GaussianNoise,
    "salt_pepper": SaltPepper,
    "paper_texture": PaperTexture,
    "scanner_artifacts": ScannerArtifacts,
}


class AugmentationPipeline:
    """
    Orchestrates multiple augmentations with proper ordering.

    Augmentations are applied in the following order:
    1. Geometric (perspective_warp) - affects bboxes
    2. Degradation (wrinkle, edge_damage, stain) - visual artifacts
    3. Lighting (lighting_variation, shadow)
    4. Texture (paper_texture, scanner_artifacts)
    5. Blur (gaussian_blur, motion_blur)
    6. Noise (gaussian_noise, salt_pepper) - applied last
    """

    STAGE_ORDER = [
        "geometric",
        "degradation",
        "lighting",
        "texture",
        "blur",
        "noise",
    ]

    STAGE_MAPPING = {
        "perspective_warp": "geometric",
        "wrinkle": "degradation",
        "edge_damage": "degradation",
        "stain": "degradation",
        "lighting_variation": "lighting",
        "shadow": "lighting",
        "paper_texture": "texture",
        "scanner_artifacts": "texture",
        "gaussian_blur": "blur",
        "motion_blur": "blur",
        "gaussian_noise": "noise",
        "salt_pepper": "noise",
    }

    def __init__(self, config: AugmentationConfig) -> None:
        """
        Initialize pipeline with configuration.

        Args:
            config: Augmentation configuration.
        """
        self.config = config
        self._rng = np.random.default_rng(config.seed)
        self._augmentations = self._build_augmentations()

    def _build_augmentations(
        self,
    ) -> list[tuple[str, BaseAugmentation, float]]:
        """Build ordered list of (name, augmentation, probability) tuples."""
        augmentations: list[tuple[str, BaseAugmentation, float]] = []

        for aug_name, aug_class in AUGMENTATION_REGISTRY.items():
            params: AugmentationParams = getattr(self.config, aug_name)
            if params.enabled:
                aug = aug_class(params.params)
                augmentations.append((aug_name, aug, params.probability))

        # Sort by stage order
        def sort_key(item: tuple[str, BaseAugmentation, float]) -> int:
            name, _, _ = item
            stage = self.STAGE_MAPPING[name]
            return self.STAGE_ORDER.index(stage)

        return sorted(augmentations, key=sort_key)

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
    ) -> AugmentationResult:
        """
        Apply augmentation pipeline to image.

        Args:
            image: Input image (H, W, C) as numpy array with dtype uint8.
            bboxes: Optional bounding boxes in YOLO format (N, 5).

        Returns:
            AugmentationResult with augmented image and optionally adjusted bboxes.
        """
        current_image = image.copy()
        current_bboxes = bboxes.copy() if bboxes is not None else None
        applied_augmentations: list[str] = []

        for name, aug, probability in self._augmentations:
            if self._rng.random() < probability:
                result = aug.apply(current_image, current_bboxes, self._rng)
                current_image = result.image
                if result.bboxes is not None and self.config.preserve_bboxes:
                    current_bboxes = result.bboxes
                applied_augmentations.append(name)

        return AugmentationResult(
            image=current_image,
            bboxes=current_bboxes,
            metadata={"applied_augmentations": applied_augmentations},
        )

    def preview(
        self,
        image: np.ndarray,
        augmentation_name: str,
    ) -> np.ndarray:
        """
        Preview a single augmentation deterministically.

        Args:
            image: Input image.
            augmentation_name: Name of augmentation to preview.

        Returns:
            Augmented image.

        Raises:
            ValueError: If augmentation_name is not recognized.
        """
        if augmentation_name not in AUGMENTATION_REGISTRY:
            raise ValueError(f"Unknown augmentation: {augmentation_name}")

        params: AugmentationParams = getattr(self.config, augmentation_name)
        aug = AUGMENTATION_REGISTRY[augmentation_name](params.params)

        # Use deterministic RNG for preview
        preview_rng = np.random.default_rng(42)
        result = aug.apply(image.copy(), rng=preview_rng)
        return result.image


def get_available_augmentations() -> list[dict[str, Any]]:
    """
    Get list of available augmentations with metadata.

    Returns:
        List of dictionaries with augmentation info.
    """
    augmentations = []
    for name, aug_class in AUGMENTATION_REGISTRY.items():
        augmentations.append({
            "name": name,
            "description": aug_class.__doc__ or "",
            "affects_geometry": aug_class.affects_geometry,
            "stage": AugmentationPipeline.STAGE_MAPPING[name],
        })
    return augmentations
