"""Augmentation service for handling augmentation operations."""

import base64
import io
import re
import uuid
from pathlib import Path
from typing import Any

import numpy as np
from fastapi import HTTPException
from PIL import Image

from backend.data.repositories import DocumentRepository, DatasetRepository
from backend.web.schemas.admin.augmentation import (
    AugmentationBatchResponse,
    AugmentationConfigSchema,
    AugmentationPreviewResponse,
    AugmentedDatasetItem,
    AugmentedDatasetListResponse,
)

# Constants
PREVIEW_MAX_SIZE = 800
PREVIEW_SEED = 42
UUID_PATTERN = re.compile(
    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",
    re.IGNORECASE,
)


class AugmentationService:
    """Service for augmentation operations."""

    def __init__(
        self,
        doc_repo: DocumentRepository | None = None,
        dataset_repo: DatasetRepository | None = None,
    ) -> None:
        """Initialize service with repository connections."""
        self.doc_repo = doc_repo or DocumentRepository()
        self.dataset_repo = dataset_repo or DatasetRepository()

    def _validate_uuid(self, value: str, field_name: str = "ID") -> None:
        """
        Validate UUID format to prevent path traversal.

        Args:
            value: Value to validate.
            field_name: Field name for error message.

        Raises:
            HTTPException: If value is not a valid UUID.
        """
        if not UUID_PATTERN.match(value):
            raise HTTPException(
                status_code=400,
                detail=f"Invalid {field_name} format: {value}",
            )

    async def preview_single(
        self,
        document_id: str,
        page: int,
        augmentation_type: str,
        params: dict[str, Any],
    ) -> AugmentationPreviewResponse:
        """
        Preview a single augmentation on a document page.

        Args:
            document_id: Document UUID.
            page: Page number (1-indexed).
            augmentation_type: Name of augmentation to apply.
            params: Override parameters.

        Returns:
            Preview response with image URLs.

        Raises:
            HTTPException: If document not found or augmentation invalid.
        """
        from shared.augmentation.config import AugmentationConfig, AugmentationParams
        from shared.augmentation.pipeline import AUGMENTATION_REGISTRY, AugmentationPipeline

        # Validate augmentation type
        if augmentation_type not in AUGMENTATION_REGISTRY:
            raise HTTPException(
                status_code=400,
                detail=f"Unknown augmentation type: {augmentation_type}. "
                f"Available: {list(AUGMENTATION_REGISTRY.keys())}",
            )

        # Get document and load image
        image = await self._load_document_page(document_id, page)

        # Create config with only this augmentation enabled
        config_kwargs = {
            augmentation_type: AugmentationParams(
                enabled=True,
                probability=1.0,  # Always apply for preview
                params=params,
            ),
            "seed": PREVIEW_SEED,  # Deterministic preview
        }
        config = AugmentationConfig(**config_kwargs)
        pipeline = AugmentationPipeline(config)

        # Apply augmentation
        result = pipeline.apply(image)

        # Convert to base64 URLs
        original_url = self._image_to_data_url(image)
        preview_url = self._image_to_data_url(result.image)

        return AugmentationPreviewResponse(
            preview_url=preview_url,
            original_url=original_url,
            applied_params=params,
        )

    async def preview_config(
        self,
        document_id: str,
        page: int,
        config: AugmentationConfigSchema,
    ) -> AugmentationPreviewResponse:
        """
        Preview full augmentation config on a document page.

        Args:
            document_id: Document UUID.
            page: Page number (1-indexed).
            config: Full augmentation configuration.

        Returns:
            Preview response with image URLs.
        """
        from shared.augmentation.config import AugmentationConfig
        from shared.augmentation.pipeline import AugmentationPipeline

        # Load image
        image = await self._load_document_page(document_id, page)

        # Convert Pydantic model to internal config
        config_dict = config.model_dump()
        internal_config = AugmentationConfig.from_dict(config_dict)
        pipeline = AugmentationPipeline(internal_config)

        # Apply augmentation
        result = pipeline.apply(image)

        # Convert to base64 URLs
        original_url = self._image_to_data_url(image)
        preview_url = self._image_to_data_url(result.image)

        return AugmentationPreviewResponse(
            preview_url=preview_url,
            original_url=original_url,
            applied_params=config_dict,
        )

    async def create_augmented_dataset(
        self,
        source_dataset_id: str,
        config: AugmentationConfigSchema,
        output_name: str,
        multiplier: int,
    ) -> AugmentationBatchResponse:
        """
        Create a new augmented dataset from an existing dataset.

        Args:
            source_dataset_id: Source dataset UUID.
            config: Augmentation configuration.
            output_name: Name for the new dataset.
            multiplier: Number of augmented copies per image.

        Returns:
            Batch response with task ID.

        Raises:
            HTTPException: If source dataset not found.
        """
        # Validate source dataset exists
        try:
            source_dataset = self.dataset_repo.get(source_dataset_id)
            if source_dataset is None:
                raise HTTPException(
                    status_code=404,
                    detail=f"Source dataset not found: {source_dataset_id}",
                )
        except Exception as e:
            raise HTTPException(
                status_code=404,
                detail=f"Source dataset not found: {source_dataset_id}",
            ) from e

        # Create task ID for background processing
        task_id = str(uuid.uuid4())

        # Estimate total images
        estimated_images = (
            source_dataset.total_images * multiplier
            if hasattr(source_dataset, "total_images")
            else 0
        )

        # TODO: Queue background task for actual augmentation
        # For now, return pending status

        return AugmentationBatchResponse(
            task_id=task_id,
            status="pending",
            message=f"Augmentation task queued for dataset '{output_name}'",
            estimated_images=estimated_images,
        )

    async def list_augmented_datasets(
        self,
        limit: int = 20,
        offset: int = 0,
    ) -> AugmentedDatasetListResponse:
        """
        List augmented datasets.

        Args:
            limit: Maximum number of datasets to return.
            offset: Number of datasets to skip.

        Returns:
            List response with datasets.
        """
        # TODO: Implement actual database query for augmented datasets
        # For now, return empty list

        return AugmentedDatasetListResponse(
            total=0,
            limit=limit,
            offset=offset,
            datasets=[],
        )

    async def _load_document_page(
        self,
        document_id: str,
        page: int,
    ) -> np.ndarray:
        """
        Load a document page as numpy array.

        Args:
            document_id: Document UUID.
            page: Page number (1-indexed).

        Returns:
            Image as numpy array (H, W, C) with dtype uint8.

        Raises:
            HTTPException: If document or page not found.
        """
        # Validate document_id format to prevent path traversal
        self._validate_uuid(document_id, "document_id")

        # Get document from database
        try:
            document = self.doc_repo.get(document_id)
            if document is None:
                raise HTTPException(
                    status_code=404,
                    detail=f"Document not found: {document_id}",
                )
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(
                status_code=404,
                detail=f"Document not found: {document_id}",
            ) from e

        # Get image path for page
        if hasattr(document, "images_dir"):
            images_dir = Path(document.images_dir)
        else:
            # Fallback to constructed path
            from backend.web.core.config import get_settings

            settings = get_settings()
            images_dir = Path(settings.admin_storage_path) / "documents" / document_id / "images"

        # Find image for page
        page_idx = page - 1  # Convert to 0-indexed
        image_files = sorted(images_dir.glob("*.png")) + sorted(images_dir.glob("*.jpg"))

        if page_idx >= len(image_files):
            raise HTTPException(
                status_code=404,
                detail=f"Page {page} not found for document {document_id}",
            )

        # Load image
        image_path = image_files[page_idx]
        pil_image = Image.open(image_path).convert("RGB")
        return np.array(pil_image)

    def _image_to_data_url(self, image: np.ndarray) -> str:
        """Convert numpy image to base64 data URL."""
        pil_image = Image.fromarray(image)

        # Resize for preview if too large
        max_size = PREVIEW_MAX_SIZE
        if max(pil_image.size) > max_size:
            ratio = max_size / max(pil_image.size)
            new_size = (int(pil_image.width * ratio), int(pil_image.height * ratio))
            pil_image = pil_image.resize(new_size, Image.Resampling.LANCZOS)

        # Convert to base64
        buffer = io.BytesIO()
        pil_image.save(buffer, format="PNG")
        base64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")

        return f"data:image/png;base64,{base64_data}"
