"""Main extraction orchestrator for transform metadata."""

from __future__ import annotations

import inspect
from typing import TYPE_CHECKING, Any, cast

if TYPE_CHECKING:
    from typing import Literal

import albumentations as A

from albu_spec.docstring_parser import DocstringParser
from albu_spec.models import ConstraintInfo, ParameterMetadata, TransformCollection, TransformMetadata
from albu_spec.schema_parser import SchemaParser
from albu_spec.type_formatters import TypeFormatter

# Transforms to ignore
IGNORED_CLASSES = {
    "Lambda",
    "BasicTransform",
    "DualTransform",
    "ImageOnlyTransform",
    "Transform3D",
    "TextImage",
    "PiecewiseAffine",
    "OverlayElements",
    "BaseTransformInitSchema",
}


class TransformMetadataExtractor:
    """Extract comprehensive metadata from Albumentations transforms."""

    def __init__(self) -> None:
        """Initialize the metadata extractor."""
        self.schema_parser = SchemaParser()
        self.docstring_parser = DocstringParser()
        self.type_formatter = TypeFormatter()

    def get_transform_metadata(self, transform_class: type) -> TransformMetadata:
        """Extract complete metadata for a single transform.

        Args:
            transform_class: The transform class to analyze

        Returns:
            TransformMetadata object containing all extracted information

        """
        # Get basic information
        name = transform_class.__name__
        module = transform_class.__module__
        transform_type = self._get_transform_type(transform_class)
        targets = self._get_targets(transform_class)
        has_init_schema = hasattr(transform_class, "InitSchema")

        # Get docstring information
        docstring = transform_class.__doc__
        docstring_short = self.docstring_parser.get_short_description(transform_class)
        param_descriptions = self.docstring_parser.parse_docstring(transform_class)

        # Get schema constraints
        schema_constraints = self.schema_parser.extract_schema_constraints(transform_class)

        # Extract parameters from __init__
        parameters = self._extract_parameters(
            transform_class,
            param_descriptions,
            schema_constraints,
        )

        # Cast transform_type to Literal type expected by TransformMetadata
        if transform_type in ("image_only", "dual", "transforms_3d"):
            valid_transform_type = cast(
                "Literal['image_only', 'dual', 'transforms_3d', 'unknown']",
                transform_type,
            )
        else:
            valid_transform_type = "unknown"

        return TransformMetadata(
            name=name,
            module=module,
            transform_type=valid_transform_type,
            targets=targets,
            parameters=parameters,
            docstring=docstring,
            docstring_short=docstring_short,
            has_init_schema=has_init_schema,
        )

    def _extract_parameters(
        self,
        transform_class: type,
        param_descriptions: dict[str, str],
        schema_constraints: dict[str, ConstraintInfo],
    ) -> dict[str, ParameterMetadata]:
        """Extract parameter metadata from __init__ signature.

        Args:
            transform_class: The transform class
            param_descriptions: Parameter descriptions from docstring
            schema_constraints: Constraints from InitSchema

        Returns:
            Dictionary mapping parameter names to their metadata

        """
        parameters: dict[str, ParameterMetadata] = {}

        try:
            # Get signature from the class __init__ method
            # Access through the class, not instance, to avoid mypy warning
            init_method = transform_class.__init__
            init_signature = inspect.signature(init_method)
        except (ValueError, TypeError, AttributeError):
            return parameters

        for param_name, param in init_signature.parameters.items():
            # Skip self and strict (strict is in InitSchema but not actually in __init__)
            if param_name in {"self", "strict"}:
                continue

            # Get type hint
            type_hint = self._format_type_hint(param.annotation, param_name, transform_class)

            # Get default value
            default_value = param.default if param.default is not inspect.Parameter.empty else None

            # Format default value
            formatted_default = self._format_default_value(default_value)

            # Get description from docstring
            description = param_descriptions.get(param_name)

            # Get constraints from schema or from type annotation
            constraints = schema_constraints.get(param_name)
            if constraints is None and param.annotation is not inspect.Parameter.empty:
                # Try to extract from Annotated types
                constraints = self.schema_parser.extract_annotated_constraints(param.annotation)

            parameters[param_name] = ParameterMetadata(
                name=param_name,
                type_hint=type_hint,
                default=formatted_default,
                description=description,
                constraints=constraints,
            )

        return parameters

    def _format_type_hint(self, annotation: object, param_name: str, transform_class: type) -> str | list[Any]:
        """Format type annotation as human-readable string.

        Args:
            annotation: Type annotation
            param_name: Parameter name (for context)
            transform_class: Transform class (for InitSchema lookup)

        Returns:
            Formatted type string or list of strings for Literal types

        """
        if annotation is inspect.Parameter.empty:
            return "Any"

        # Try to get more precise type from InitSchema if available
        if hasattr(transform_class, "InitSchema"):
            init_schema = transform_class.InitSchema
            if hasattr(init_schema, "model_fields") and param_name in init_schema.model_fields:
                field_info = init_schema.model_fields[param_name]
                if hasattr(field_info, "annotation"):
                    annotation = field_info.annotation

        return self._format_type(annotation)

    def _format_type(self, type_annotation: object) -> str | list[Any]:
        """Format a type annotation into a readable string.

        Args:
            type_annotation: Type annotation to format

        Returns:
            Formatted type string or list for Literal types

        """
        return self.type_formatter.format(type_annotation)

    def _format_default_value(self, value: object) -> Any:
        """Format default value for display.

        Args:
            value: Default value to format

        Returns:
            Formatted default value

        """
        if callable(value) and not isinstance(value, type) and hasattr(value, "__name__"):
            return f"<function {value.__name__}>"

        return value

    def _get_transform_type(self, transform_class: type) -> str:
        """Determine the type of transform.

        Args:
            transform_class: Transform class to analyze

        Returns:
            Transform type string

        """
        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"

        return "unknown"

    def _get_targets(self, transform_class: type) -> list[str]:
        """Get supported targets for the transform.

        Args:
            transform_class: Transform class to analyze

        Returns:
            List of target names

        """
        targets: list[str] = []

        if hasattr(transform_class, "_targets"):
            targets_attr = transform_class._targets

            # Handle various types of _targets
            if isinstance(targets_attr, (list, tuple)) or (
                hasattr(targets_attr, "__iter__") and not isinstance(targets_attr, str)
            ):
                for target in targets_attr:
                    target_str = getattr(target, "value", target)
                    if isinstance(target_str, str):
                        targets.append(target_str.lower())

        return targets

    def get_all_transforms_metadata(self) -> TransformCollection:
        """Extract metadata for all Albumentations transforms.

        Returns:
            TransformCollection with all transforms grouped by type

        """
        collection = TransformCollection()

        # Find all transform classes
        for name, obj in inspect.getmembers(A, predicate=inspect.isclass):
            # Skip ignored classes
            if name in IGNORED_CLASSES:
                continue

            # Check if it's a transform
            try:
                if not issubclass(obj, A.BasicTransform):
                    continue
                if obj is A.BasicTransform:
                    continue
            except TypeError:
                continue

            # Extract metadata
            try:
                metadata = self.get_transform_metadata(obj)

                # Add to appropriate category
                if metadata.transform_type == "image_only":
                    collection.image_only.append(metadata)
                elif metadata.transform_type == "dual":
                    collection.dual.append(metadata)
                elif metadata.transform_type == "transforms_3d":
                    collection.transforms_3d.append(metadata)
                else:
                    collection.unknown.append(metadata)
            except (ValueError, TypeError, AttributeError):
                # Skip transforms that fail to extract
                continue

        return collection


# Public API functions
def get_transform_metadata(transform_class: type) -> TransformMetadata:
    """Extract metadata for a single transform.

    Args:
        transform_class: The transform class to analyze

    Returns:
        TransformMetadata object with all extracted information

    Example:
        >>> import albumentations as A
        >>> from albu_spec import get_transform_metadata
        >>> metadata = get_transform_metadata(A.Blur)
        >>> print(metadata.name)
        'Blur'
        >>> print(metadata.parameters['blur_limit'].type_hint)
        'tuple[int, int] | int'

    """
    extractor = TransformMetadataExtractor()
    return extractor.get_transform_metadata(transform_class)


def get_all_transforms_metadata() -> TransformCollection:
    """Extract metadata for all Albumentations transforms.

    Returns:
        TransformCollection with all transforms grouped by type

    Example:
        >>> from albu_spec import get_all_transforms_metadata
        >>> collection = get_all_transforms_metadata()
        >>> print(f"Found {collection.total_count} transforms")
        >>> print(f"Image-only: {len(collection.image_only)}")
        >>> print(f"Dual: {len(collection.dual)}")

    """
    extractor = TransformMetadataExtractor()
    return extractor.get_all_transforms_metadata()
