"""
Texture augmentation transforms.

Provides texture effects for document image augmentation:
- PaperTexture: Adds paper grain/texture
- ScannerArtifacts: Adds scanner line and dust artifacts
"""

import cv2
import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class PaperTexture(BaseAugmentation):
    """
    Adds paper texture/grain to the image.

    Simulates different paper types and ages.

    Parameters:
        texture_type: Type of texture ("random", "fine", "coarse") (default: "random").
        intensity: Texture intensity, float or (min, max) tuple (default: (0.05, 0.15)).
    """

    name = "paper_texture"
    affects_geometry = False

    def _validate_params(self) -> None:
        intensity = self.params.get("intensity", (0.05, 0.15))
        if isinstance(intensity, (int, float)):
            if not (0 < intensity <= 1):
                raise ValueError("intensity must be between 0 and 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()

        h, w = image.shape[:2]
        texture_type = self.params.get("texture_type", "random")
        intensity = self.params.get("intensity", (0.05, 0.15))

        if texture_type == "random":
            texture_type = rng.choice(["fine", "coarse"])

        if isinstance(intensity, tuple):
            intensity = rng.uniform(intensity[0], intensity[1])

        # Generate base noise
        if texture_type == "fine":
            # Fine grain texture
            noise = rng.uniform(-1, 1, (h, w)).astype(np.float32)
            noise = cv2.GaussianBlur(noise, (3, 3), 0)
        else:
            # Coarse texture
            # Generate at lower resolution and upscale
            small_h, small_w = h // 4, w // 4
            noise = rng.uniform(-1, 1, (small_h, small_w)).astype(np.float32)
            noise = cv2.resize(noise, (w, h), interpolation=cv2.INTER_LINEAR)
            noise = cv2.GaussianBlur(noise, (5, 5), 0)

        # Apply texture
        output = image.astype(np.float32)
        noise_3d = noise[:, :, np.newaxis] * intensity * 255
        output = output + noise_3d

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

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

    def get_preview_params(self) -> dict:
        return {"texture_type": "coarse", "intensity": 0.15}


class ScannerArtifacts(BaseAugmentation):
    """
    Adds scanner artifacts to the image.

    Simulates scanner imperfections like lines and dust spots.

    Parameters:
        line_probability: Probability of adding scan lines (default: 0.3).
        dust_probability: Probability of adding dust spots (default: 0.4).
    """

    name = "scanner_artifacts"
    affects_geometry = False

    def _validate_params(self) -> None:
        line_prob = self.params.get("line_probability", 0.3)
        dust_prob = self.params.get("dust_probability", 0.4)
        if not (0 <= line_prob <= 1):
            raise ValueError("line_probability must be between 0 and 1")
        if not (0 <= dust_prob <= 1):
            raise ValueError("dust_probability must be between 0 and 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()

        h, w = image.shape[:2]
        line_probability = self.params.get("line_probability", 0.3)
        dust_probability = self.params.get("dust_probability", 0.4)

        output = image.copy()

        # Add scan lines
        if rng.random() < line_probability:
            num_lines = rng.integers(1, 4)
            for _ in range(num_lines):
                y = rng.integers(0, h)
                thickness = rng.integers(1, 3)
                # Light or dark line
                color = rng.integers(200, 240) if rng.random() > 0.5 else rng.integers(50, 100)

                # Make line partially transparent
                alpha = rng.uniform(0.3, 0.6)
                for dy in range(thickness):
                    if y + dy < h:
                        output[y + dy, :] = (
                            output[y + dy, :].astype(np.float32) * (1 - alpha) +
                            color * alpha
                        ).astype(np.uint8)

        # Add dust spots
        if rng.random() < dust_probability:
            num_dust = rng.integers(5, 20)
            for _ in range(num_dust):
                x = rng.integers(0, w)
                y = rng.integers(0, h)
                radius = rng.integers(1, 3)

                # Dark dust spot
                color = rng.integers(50, 120)
                cv2.circle(output, (x, y), radius, int(color), -1)

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

    def get_preview_params(self) -> dict:
        return {"line_probability": 0.8, "dust_probability": 0.8}
