"""
Lighting augmentation transforms.

Provides lighting effects for document image augmentation:
- LightingVariation: Adjusts brightness and contrast
- Shadow: Adds shadow overlay effects
"""

import cv2
import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class LightingVariation(BaseAugmentation):
    """
    Adjusts image brightness and contrast.

    Simulates different lighting conditions during document capture.
    Safe for documents with conservative default parameters.

    Parameters:
        brightness_range: (min, max) brightness adjustment (default: (-0.1, 0.1)).
        contrast_range: (min, max) contrast multiplier (default: (0.9, 1.1)).
    """

    name = "lighting_variation"
    affects_geometry = False

    def _validate_params(self) -> None:
        brightness = self.params.get("brightness_range", (-0.1, 0.1))
        contrast = self.params.get("contrast_range", (0.9, 1.1))

        if not isinstance(brightness, tuple) or len(brightness) != 2:
            raise ValueError("brightness_range must be a (min, max) tuple")
        if not isinstance(contrast, tuple) or len(contrast) != 2:
            raise ValueError("contrast_range must be a (min, max) tuple")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        brightness_range = self.params.get("brightness_range", (-0.1, 0.1))
        contrast_range = self.params.get("contrast_range", (0.9, 1.1))

        # Random brightness and contrast
        brightness = rng.uniform(brightness_range[0], brightness_range[1])
        contrast = rng.uniform(contrast_range[0], contrast_range[1])

        # Apply adjustments
        adjusted = image.astype(np.float32)

        # Contrast adjustment (multiply around mean)
        mean = adjusted.mean()
        adjusted = (adjusted - mean) * contrast + mean

        # Brightness adjustment (add offset)
        adjusted = adjusted + brightness * 255

        # Clip and convert back
        adjusted = np.clip(adjusted, 0, 255).astype(np.uint8)

        return AugmentationResult(
            image=adjusted,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"brightness": brightness, "contrast": contrast},
        )

    def get_preview_params(self) -> dict:
        return {"brightness_range": (-0.15, 0.15), "contrast_range": (0.85, 1.15)}


class Shadow(BaseAugmentation):
    """
    Adds shadow overlay effects to the image.

    Simulates shadows from objects or hands during document capture.

    Parameters:
        num_shadows: Number of shadow regions, int or (min, max) tuple (default: (1, 2)).
        opacity: Shadow darkness, float or (min, max) tuple (default: (0.2, 0.4)).
    """

    name = "shadow"
    affects_geometry = False

    def _validate_params(self) -> None:
        opacity = self.params.get("opacity", (0.2, 0.4))
        if isinstance(opacity, (int, float)):
            if not (0 <= opacity <= 1):
                raise ValueError("opacity must be between 0 and 1")
        elif isinstance(opacity, tuple):
            if not (0 <= opacity[0] <= opacity[1] <= 1):
                raise ValueError("opacity tuple must be in range [0, 1]")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        num_shadows = self.params.get("num_shadows", (1, 2))
        opacity = self.params.get("opacity", (0.2, 0.4))

        if isinstance(num_shadows, tuple):
            num_shadows = rng.integers(num_shadows[0], num_shadows[1] + 1)
        if isinstance(opacity, tuple):
            opacity = rng.uniform(opacity[0], opacity[1])

        h, w = image.shape[:2]
        output = image.astype(np.float32)

        for _ in range(num_shadows):
            # Generate random shadow polygon
            num_vertices = rng.integers(3, 6)
            vertices = []

            # Start from a random edge
            edge = rng.integers(0, 4)
            if edge == 0:  # Top
                start = (rng.integers(0, w), 0)
            elif edge == 1:  # Right
                start = (w, rng.integers(0, h))
            elif edge == 2:  # Bottom
                start = (rng.integers(0, w), h)
            else:  # Left
                start = (0, rng.integers(0, h))

            vertices.append(start)

            # Add random vertices
            for _ in range(num_vertices - 1):
                x = rng.integers(0, w)
                y = rng.integers(0, h)
                vertices.append((x, y))

            # Create shadow mask
            mask = np.zeros((h, w), dtype=np.float32)
            pts = np.array(vertices, dtype=np.int32).reshape((-1, 1, 2))
            cv2.fillPoly(mask, [pts], 1.0)

            # Blur the mask for soft edges
            blur_size = max(31, min(h, w) // 10)
            if blur_size % 2 == 0:
                blur_size += 1
            mask = cv2.GaussianBlur(mask, (blur_size, blur_size), 0)

            # Apply shadow
            shadow_factor = 1 - opacity * mask[:, :, np.newaxis]
            output = output * shadow_factor

        output = np.clip(output, 0, 255).astype(np.uint8)

        return AugmentationResult(
            image=output,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"num_shadows": num_shadows, "opacity": opacity},
        )

    def get_preview_params(self) -> dict:
        return {"num_shadows": 1, "opacity": 0.3}
