"""
Dataset Augmenter Module.

Applies augmentation pipeline to YOLO datasets,
creating new augmented images and label files.
"""

import logging
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from shared.augmentation.config import AugmentationConfig, AugmentationParams
from shared.augmentation.pipeline import AugmentationPipeline

logger = logging.getLogger(__name__)


class DatasetAugmenter:
    """
    Augments YOLO datasets by creating new images and label files.

    Reads images from dataset/images/train/ and labels from dataset/labels/train/,
    applies augmentation pipeline, and saves augmented versions with "_augN" suffix.
    """

    def __init__(
        self,
        config: dict[str, Any],
        seed: int | None = None,
    ) -> None:
        """
        Initialize augmenter with configuration.

        Args:
            config: Dictionary mapping augmentation names to their settings.
                    Each augmentation should have 'enabled', 'probability', and 'params'.
            seed: Random seed for reproducibility.
        """
        self._config_dict = config
        self._seed = seed
        self._config = self._build_config(config, seed)

    def _build_config(
        self,
        config_dict: dict[str, Any],
        seed: int | None,
    ) -> AugmentationConfig:
        """Build AugmentationConfig from dictionary."""
        kwargs: dict[str, Any] = {"seed": seed, "preserve_bboxes": True}

        for aug_name, aug_settings in config_dict.items():
            if aug_name in AugmentationConfig._AUGMENTATION_FIELDS:
                kwargs[aug_name] = AugmentationParams(
                    enabled=aug_settings.get("enabled", False),
                    probability=aug_settings.get("probability", 0.5),
                    params=aug_settings.get("params", {}),
                )

        return AugmentationConfig(**kwargs)

    def augment_dataset(
        self,
        dataset_path: Path,
        multiplier: int = 1,
        split: str = "train",
    ) -> dict[str, int]:
        """
        Augment a YOLO dataset.

        Args:
            dataset_path: Path to dataset root (containing images/ and labels/).
            multiplier: Number of augmented copies per original image.
            split: Which split to augment (default: "train").

        Returns:
            Summary dict with original_images, augmented_images, total_images.
        """
        images_dir = dataset_path / "images" / split
        labels_dir = dataset_path / "labels" / split

        if not images_dir.exists():
            raise ValueError(f"Images directory not found: {images_dir}")

        # Find all images
        image_extensions = ("*.png", "*.jpg", "*.jpeg")
        image_files: list[Path] = []
        for ext in image_extensions:
            image_files.extend(images_dir.glob(ext))

        original_count = len(image_files)
        augmented_count = 0

        if multiplier <= 0:
            return {
                "original_images": original_count,
                "augmented_images": 0,
                "total_images": original_count,
            }

        # Process each image
        for img_path in image_files:
            # Load image
            pil_image = Image.open(img_path).convert("RGB")
            image = np.array(pil_image)

            # Load corresponding label
            label_path = labels_dir / f"{img_path.stem}.txt"
            bboxes = self._load_bboxes(label_path) if label_path.exists() else None

            # Create multiple augmented versions
            for aug_idx in range(multiplier):
                # Create pipeline with adjusted seed for each augmentation
                aug_seed = None
                if self._seed is not None:
                    aug_seed = self._seed + aug_idx + hash(img_path.stem) % 10000

                pipeline = AugmentationPipeline(
                    self._build_config(self._config_dict, aug_seed)
                )

                # Apply augmentation
                result = pipeline.apply(image, bboxes)

                # Save augmented image
                aug_name = f"{img_path.stem}_aug{aug_idx}{img_path.suffix}"
                aug_img_path = images_dir / aug_name
                aug_pil = Image.fromarray(result.image)
                aug_pil.save(aug_img_path)

                # Save augmented label
                aug_label_path = labels_dir / f"{img_path.stem}_aug{aug_idx}.txt"
                self._save_bboxes(aug_label_path, result.bboxes)

                augmented_count += 1

        logger.info(
            "Dataset augmentation complete: %d original, %d augmented",
            original_count,
            augmented_count,
        )

        return {
            "original_images": original_count,
            "augmented_images": augmented_count,
            "total_images": original_count + augmented_count,
        }

    def _load_bboxes(self, label_path: Path) -> np.ndarray | None:
        """
        Load bounding boxes from YOLO label file.

        Args:
            label_path: Path to label file.

        Returns:
            Array of shape (N, 5) with class_id, x_center, y_center, width, height.
            Returns None if file is empty or doesn't exist.
        """
        if not label_path.exists():
            return None

        content = label_path.read_text().strip()
        if not content:
            return None

        bboxes = []
        for line in content.split("\n"):
            parts = line.strip().split()
            if len(parts) == 5:
                class_id = int(parts[0])
                x_center = float(parts[1])
                y_center = float(parts[2])
                width = float(parts[3])
                height = float(parts[4])
                bboxes.append([class_id, x_center, y_center, width, height])

        if not bboxes:
            return None

        return np.array(bboxes, dtype=np.float32)

    def _save_bboxes(self, label_path: Path, bboxes: np.ndarray | None) -> None:
        """
        Save bounding boxes to YOLO label file.

        Args:
            label_path: Path to save label file.
            bboxes: Array of shape (N, 5) or None for empty labels.
        """
        if bboxes is None or len(bboxes) == 0:
            label_path.write_text("")
            return

        lines = []
        for bbox in bboxes:
            class_id = int(bbox[0])
            x_center = bbox[1]
            y_center = bbox[2]
            width = bbox[3]
            height = bbox[4]
            lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")

        label_path.write_text("\n".join(lines))
