import pytest
import albumentations as A


from src.transforms.info import get_transform_info
from src.transforms.constants import CV2ConstantHandler

# Initialize utilities needed for tests
cv2_handler = CV2ConstantHandler()

INTERPOLATION_NAMES = cv2_handler.interpolation_names
BORDER_NAMES = cv2_handler.border_names


def normalize_type_str(type_str: str) -> set[str]:
    """
    Convert a type string into a set of components for order-independent comparison.

    Examples:
        "int | str | None" -> {"int", "str", "None"}
        "float | tuple[float, ...]" -> {"float", "tuple[float, ...]"}
        "dict[str, float | tuple[float, float]]" -> {"dict[str, float | tuple[float, float]]"}
    """
    if not isinstance(type_str, str):
        return {str(x) for x in type_str}

    # If it's not a union type, return as a single-element set
    if " | " not in type_str:
        return {type_str}

    # Split union types and return as set
    return set(part.strip() for part in type_str.split(" | "))


@pytest.mark.parametrize("transform_class,expected", [
    (A.Normalize, {
        "name": "Normalize",
        "type": "image_only",
        "targets": ["image", "volume"],
        "fields": {
            "mean": {
                "type": "float | tuple[float, ...] | None",
                "default": (0.485, 0.456, 0.406),
            },
            "std": {
                "type": "float | tuple[float, ...] | None",
                "default": (0.229, 0.224, 0.225),
            },
            "max_pixel_value": {
                "type": "float | None",
                "default": 255.0,
            },
            "normalization": {
                "type": ["standard", "image", "image_per_channel", "min_max", "min_max_per_channel"],
                "default": "standard",
            },
        },
    }),
    (A.AdditiveNoise, {
        "name": "AdditiveNoise",
        "type": "image_only",
        "targets": ["image", "volume"],
        "fields": {
            "noise_type": {
                "type": ["uniform", "gaussian", "laplace", "beta"],
                "default": "uniform",
            },
            "spatial_mode": {
                "type": ["constant", "per_pixel", "shared"],
                "default": "constant",
            },
            "approximation": {
                "type": "float",
                "default": 1.0,
            },
            "noise_params": {
                "type": "dict[str, Any] | None",
                "default": None,
            }
        }
    }),
    (A.FDA, {
        "name": "FDA",
        "type": "image_only",
        "targets": ["image", "volume"],
        "fields": {
            "reference_images": {
                "type": "tuple[Any, ...] | None",
                "default": None,
            },
            "beta_limit": {
                "type": "float | tuple[float, float]",
                "default": (0, 0.1),
            },
            "read_fn": {
                "type": "(x: Any) => np.ndarray | None",
                "default": "<function read_rgb_image>",
            }
        }
    }),
    (A.Affine, {
        "name": "Affine",
        "type": "dual",
        "targets": ["image", "mask", "bboxes", "keypoints", "volume", "mask3d"],
        "fields": {
            "scale": {
                "type": "float | tuple[float, float] | dict[str, float | tuple[float, float]]",
                "default": (1, 1),
            },
            "translate_percent": {
                "type": "float | tuple[float, float] | dict[str, float | tuple[float, float]] | None",
                "default": None,
            },
            "translate_px": {
                "type": "tuple[int, int] | int | dict[str, int | tuple[int, int]] | None",
                "default": None,
            },
            "rotate": {
                "type": "float | tuple[float, float]",
                "default": 0,
            },
            "shear": {
                "type": "float | tuple[float, float] | dict[str, float | tuple[float, float]]",
                "default": (0, 0),
            },
            "interpolation": {
                "type": ["cv2.INTER_NEAREST",
            "cv2.INTER_LINEAR",
            "cv2.INTER_CUBIC",
            "cv2.INTER_AREA",
            "cv2.INTER_LANCZOS4",
        ],
                "default": "cv2.INTER_LINEAR",  # Test expects string "cv2.INTER_LINEAR"
            },
            "mask_interpolation": {
                "type": ["cv2.INTER_NEAREST",
                    "cv2.INTER_LINEAR",
                    "cv2.INTER_CUBIC",
                    "cv2.INTER_AREA",
                    "cv2.INTER_LANCZOS4",
                ],
                "default": "cv2.INTER_NEAREST",  # Test expects string "cv2.INTER_NEAREST"
            },
            "fill": {
                "type": "float | tuple[float, ...]",
                "default": 0,
            },
            "fill_mask": {
                "type": "float | tuple[float, ...]",
                "default": 0,
            },
            "border_mode": {
                "type": ["cv2.BORDER_CONSTANT",
                    "cv2.BORDER_REPLICATE",
                    "cv2.BORDER_REFLECT",
                    "cv2.BORDER_WRAP",
                    "cv2.BORDER_REFLECT_101",
                ],
                "default": "cv2.BORDER_CONSTANT",
            },
            "fit_output": {
                "type": "bool",
                "default": False,
            },
            "keep_ratio": {
                "type": "bool",
                "default": False,
            },
            "rotate_method": {
                "type": ["largest_box", "ellipse"],
                "default": "largest_box",
            },
            "balanced_scale": {
                "type": "bool",
                "default": False,
            },
        }
    }),
    (A.Equalize, {
        "name": "Equalize",
        "type": "image_only",
        "targets": ["image", "volume"],
        "fields": {
            "mode": {
                "type": ["cv", "pil"],
                "default": "cv",
            },
            "by_channels": {
                "type": "bool",
                "default": True,
            },
            "mask": {
                "type": "np.ndarray | (x: Any) => Any | None",
                "default": None,
            },
            "mask_params": {
                "type": "tuple[str, ...]",
                "default": (),
            },
        }
    })
])
def test_transform_info(transform_class, expected):
    """Test that get_transform_info correctly extracts information from transforms."""
    result = get_transform_info(transform_class)

    assert result.name == expected["name"]
    assert result.transform_type == expected["type"]
    assert result.targets == expected["targets"]
    assert result.parsed_docstring is not None

    # Check fields
    for key, expected_field in expected["fields"].items():
        assert key in result.fields

        # Compare types using normalize_type_str for union types
        if isinstance(expected_field["type"], str) and " | " in expected_field["type"]:
            assert normalize_type_str(result.fields[key]["type"]) == normalize_type_str(expected_field["type"])
        else:
            # For lists (like in interpolation options), keep exact order
            assert result.fields[key]["type"] == expected_field["type"]

        assert result.fields[key]["default"] == expected_field["default"]
