"""Utilities for extracting metadata and information about Albumentations transforms."""

import logging
from typing import Any, get_args

import albumentations as A
from google_docstring_parser import parse_google_docstring

from src.models.transform import Targets, TransformInfo, TransformType

from .fields import FieldExtractor

logger = logging.getLogger(__name__)  # Initialize logger before try/except

literal_type = get_args(Targets)[0]  # Get the Literal[...] from inside Sequence
literal_args = get_args(literal_type)  # Get the actual values from the Literal
ALLOWED_TARGETS = set(literal_args)


class TransformMetadataExtractor:
    """Extracts and processes metadata from Albumentations transform classes."""

    def __init__(self) -> None:
        self.field_extractor = FieldExtractor()

    def get_transform_info(self, transform_class: type[Any]) -> TransformInfo:
        """Extract complete information about a given Albumentations transform class.

        Args:
            transform_class: The transform class to analyze

        Returns:
            TransformInfo: TransformInfo containing complete transform information
        """
        # Ensure targets are extracted first to handle potential errors early
        targets = self._get_targets(transform_class)
        transform_type = self._get_transform_type(transform_class)

        # Parse the docstring if available
        if transform_class.__doc__:
            parsed_docstring = parse_google_docstring(transform_class.__doc__)
        else:
            logger.warning(f"No docstring found for {transform_class.__name__}")
            parsed_docstring = None

        return TransformInfo(
            name=transform_class.__name__,
            transform_type=transform_type,
            targets=targets,
            parsed_docstring=parsed_docstring,
            fields=self.field_extractor.extract_transform_info(transform_class),
        )

    def _get_transform_type(self, transform_class: type[Any]) -> TransformType:
        """Determine the type of transform based on its class hierarchy.

        Args:
            transform_class: The transform class to analyze

        Returns:
            TransformType: The type of the transform

        Raises:
            ValueError: If the transform type cannot be determined
        """
        if issubclass(transform_class, A.Transform3D):
            return "transforms_3d"
        if issubclass(transform_class, A.ImageOnlyTransform):
            return "image_only"
        if issubclass(transform_class, A.DualTransform):
            return "dual"
        # Fallback type
        logger.warning(
            f"Unknown transform type for class {transform_class.__name__}. "
            f"Assigning 'unknown'. Should inherit from Transform3D, ImageOnlyTransform, or DualTransform.",
        )
        return "unknown"  # Return unknown instead of raising error

    def _get_targets(self, transform_class: type[Any]) -> Targets:
        """Get the targets that a transform can be applied to.

        Args:
            transform_class: The transform class to inspect

        Returns:
            Targets: List of lowercase target names conforming to the Targets Literal type.
        """
        targets_from_class = getattr(transform_class, "_targets", {})
        raw_targets: list[Any] = []

        # Handle various ways targets might be defined (list, tuple, dict_keys, etc.)
        if isinstance(targets_from_class, (list, tuple)):
            raw_targets = [getattr(target, "value", target) for target in targets_from_class]
        elif hasattr(targets_from_class, "__iter__") and not isinstance(targets_from_class, str):
            # Convert iterables like dict_keys to list
            raw_targets = [getattr(target, "value", target) for target in targets_from_class]
        elif targets_from_class is not None:  # Handle single non-iterable target
            raw_targets = [getattr(targets_from_class, "value", targets_from_class)]
        # If targets_from_class is None or empty, raw_targets remains []

        # Filter and validate targets against the allowed Literal values
        valid_targets: list[Any] = []
        for target in raw_targets:
            if isinstance(target, str):
                lower_target = target.lower()
                if lower_target in ALLOWED_TARGETS:
                    # Append the validated string, it conforms to Literal constraint
                    valid_targets.append(lower_target)
                else:
                    logger.warning(f"Skipping unknown target '{target}' for transform {transform_class.__name__}")
            else:
                logger.warning(f"Skipping non-string target '{target}' for transform {transform_class.__name__}")

        return valid_targets


# Public interface
def get_transform_info(transform_class: type[Any]) -> TransformInfo:
    """Extract complete information about an Albumentations transform class.

    Args:
        transform_class: The transform class to analyze

    Returns:
        TransformInfo: TransformInfo object containing complete transform information

    Example:
        >>> import albumentations as A
        >>> info = get_transform_info(A.Blur)
        >>> print(info.name)
        'Blur'
        >>> print(info.transform_type)
        'image_only'
    """
    extractor = TransformMetadataExtractor()
    return extractor.get_transform_info(transform_class)
