import cv2
import pytest
from typing import Any, Callable, Literal, Sequence, Union, get_args
import numpy as np

from src.transforms.fields import FieldExtractor
from src.transforms.constants import CV2ConstantHandler

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

cv2_handler = CV2ConstantHandler()

@pytest.fixture
def field_extractor():
    return FieldExtractor()

# 1. Type Formatting Tests
@pytest.mark.parametrize("type_input,expected", [
    # Basic types
    (int, "int"),
    (float, "float"),
    (bool, "bool"),
    (str, "str"),
    (None, "None"),
    (type(None), "None"),
    (Any, "Any"),

    # Union types
    (int | None, "int | None"),
    (float | int, "float | int"),
    (str | int | None, "str | int | None"),
    (Union[int, None], "int | None"),

    # Generic types
    (list[int], "tuple[int, ...]"),
    (dict[str, int], "dict[str, int]"),
    (Sequence[int], "tuple[int, ...]"),
    (tuple[int, str], "tuple[int, str]"),
    (tuple[int, ...], "tuple[int, ...]"),

    # Callable types
    (Callable[[int], str], "(x: int) => str"),
    (Callable[[int, str], bool], "(x: int | str) => bool"),

    # Literal types
    (Literal["a", "b"], ["a", "b"]),
    (Literal[1, 2, 3], ["1", "2", "3"]),

    # Complex nested types
    (dict[str, list[int]], "dict[str, tuple[int, ...]]"),
    (dict[str, tuple[int, ...] | None], "dict[str, tuple[int, ...] | None]"),
])
def test_format_type(field_extractor, type_input, expected):
    result = field_extractor._format_type(type_input)
    assert result == expected


# 4. Default Value Tests
@pytest.mark.parametrize("value,field_name,expected", [
    (lambda x: x, None, "<function <lambda>>"),
    (None, None, None),
    (0, "border_mode", "cv2.BORDER_CONSTANT"),
    (1, "interpolation", "cv2.INTER_LINEAR"),
    (1, "mask_interpolation", "cv2.INTER_LINEAR"),
    (42, None, 42),
    ("test", None, "test"),
    (True, None, True),
    (999, "border_mode", 999),  # Invalid CV2 constant
    (-1, "interpolation", -1),  # Invalid CV2 constant
])
def test_format_default_value(field_extractor, value, field_name, expected):
    result = field_extractor.format_default_value(value, field_name)
    assert result == expected

# 5. Init Defaults Test
class SampleClass:
    def __init__(self, param1: int = 1, param2: str = "test", p: float = 0.5, strict: bool = True):
        pass

def test_extract_init_defaults(field_extractor):
    result = field_extractor.extract_init_defaults(SampleClass)
    assert result == {"param1": 1, "param2": "test"}



@pytest.mark.parametrize("callable_type,expected", [
    (
        Callable[[int], str],
        "(x: int) => str"
    ),
    (
        Callable[[int, str], bool],
        "(x: int | str) => bool"
    ),
])
def test_format_callable(field_extractor, callable_type, expected):
    result = field_extractor._format_callable(get_args(callable_type))
    assert result == expected


@pytest.mark.parametrize("type_input,field_name,expected", [
    # Basic types as strings
    ("int", "regular_field", "int"),
    ("float", "regular_field", "float"),
    ("str", "regular_field", "str"),

    # CV2 special fields with string type annotations
    (Literal[cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4], "interpolation", ['cv2.INTER_NEAREST', 'cv2.INTER_LINEAR', 'cv2.INTER_CUBIC', 'cv2.INTER_AREA', 'cv2.INTER_LANCZOS4']),
    (Literal[cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4], "mask_interpolation", ['cv2.INTER_NEAREST', 'cv2.INTER_LINEAR', 'cv2.INTER_CUBIC', 'cv2.INTER_AREA', 'cv2.INTER_LANCZOS4']),
    (Literal[cv2.BORDER_CONSTANT, cv2.BORDER_REPLICATE, cv2.BORDER_REFLECT, cv2.BORDER_WRAP, cv2.BORDER_REFLECT_101], "border_mode", ['cv2.BORDER_CONSTANT', 'cv2.BORDER_REPLICATE', 'cv2.BORDER_REFLECT', 'cv2.BORDER_WRAP', 'cv2.BORDER_REFLECT_101']),

    # Literal types as strings
    ("Literal['gaussian', 'uniform']", "noise_type", ["gaussian", "uniform"]),
    ("Literal['standard', 'simple']", "normalization", ["standard", "simple"]),

    # Sequence types as strings
    ("Sequence[Any]", "regular_field", "tuple[Any, ...]"),
    ("Sequence[int]", "regular_field", "tuple[int, ...]"),

    # Regular types for comparison
    (int, "regular_field", "int"),
    (Literal[cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4], "interpolation", ['cv2.INTER_NEAREST', 'cv2.INTER_LINEAR', 'cv2.INTER_CUBIC', 'cv2.INTER_AREA', 'cv2.INTER_LANCZOS4']),

    # Callable types with ellipsis
    (Callable[..., Any], "mask", "(x: Any) => Any"),
    (Callable[..., int], "mask", "(x: Any) => int"),

    # Complex Callable types
    (Callable[[int], Any], "mask", "(x: int) => Any"),
    (Callable[[int, str], Any], "mask", "(x: int | str) => Any"),

    # Nested types with Callable
    (Union[Callable[..., Any], None], "mask", "(x: Any) => Any | None"),
    (Union[np.ndarray, Callable[..., Any]], "mask", "np.ndarray | (x: Any) => Any"),

    # Additional complex types
    (Union[np.ndarray, Callable[..., Any], None], "mask", "np.ndarray | (x: Any) => Any | None"),
    (Union[Sequence[int], Callable[..., Any]], "mask", "tuple[int, ...] | (x: Any) => Any"),
])
def test_get_field_type_with_string_annotations(field_extractor, type_input, field_name, expected):
    """Test field type handling with both string and regular type annotations."""
    result = field_extractor._get_field_type(type_input, field_name)
    if isinstance(expected, list):
        assert result == expected
    else:
        assert normalize_type_str(result) == normalize_type_str(expected)
