import pytest
import numpy as np
import albumentations as A
from src.services.transform_service import TransformService
from src.exceptions import TransformNotFoundError
from src.models.transform import TransformInfo, TransformExecutionConfig
from src.services.image_service import ImageService
import cv2
from src.config import settings

image_service = ImageService()

# Helper to load original numpy image
def load_original_image_np(transform_name: str) -> np.ndarray:
    image_set = settings.get_image_set(transform_name)
    image_path = settings.get_image_path(image_set)
    return image_service.load_image(str(image_path))

@pytest.fixture
def transform_service():
    return TransformService()

@pytest.mark.parametrize("transform_name, expected_class", [
    ("Blur", A.Blur),
    ("Rotate", A.Rotate),
    ("HorizontalFlip", A.HorizontalFlip),
])
def test_get_transform_class_valid(transform_service, transform_name, expected_class):
    transform_class = transform_service.get_transform_class(transform_name)
    assert transform_class == expected_class

@pytest.mark.parametrize("invalid_transform", [
    "NonExistentTransform",
    "InvalidName",
    "",
])
def test_get_transform_class_invalid(transform_service, invalid_transform):
    with pytest.raises(ValueError, match=f"Transform {invalid_transform} not found"):
        transform_service.get_transform_class(invalid_transform)

@pytest.mark.parametrize("transform_name, expected_targets", [
    ("Blur", ["image", "volume"]),
    ("HorizontalFlip", ["image", "mask", "bboxes", "keypoints", "volume", "mask3d"]),
])
def test_get_transform_info(transform_service, transform_name, expected_targets):
    info = transform_service.get_transform_info(transform_name)
    assert isinstance(info, TransformInfo)
    assert info.name == transform_name
    assert sorted(info.targets) == sorted(expected_targets)

def test_get_transform_info_invalid(transform_service):
    with pytest.raises(TransformNotFoundError):
        transform_service.get_transform_info("NonExistentTransform")

def test_get_all_transforms_info(transform_service):
    result = transform_service.get_all_transforms_info()

    assert isinstance(result, dict)
    assert set(result.keys()) == {"image_only", "dual", "transforms_3d"}

    # Check that we have some transforms in each category
    assert len(result["image_only"]) > 0
    assert len(result["dual"]) > 0

    # Verify structure of returned transforms
    for category in result.values():
        for transform in category:
            assert isinstance(transform, TransformInfo)
            assert transform.name
            assert isinstance(transform.targets, list)
            assert len(transform.targets) > 0

@pytest.mark.parametrize("transform_class, params, expected_shape", [
    (A.Blur, {"blur_limit": 3}, (733, 484, 3)),
    (A.Rotate, {"limit": 45}, (733, 484, 3)),
    (A.Resize, {"height": 50, "width": 50}, (50, 50, 3)),
])
def test_execute_transform_basic(transform_service, transform_class, params, expected_shape):
    # Create a default config
    config = TransformExecutionConfig(
        targets=["image"],
        bbox_params=None,
        keypoint_params=None,
    )
    result = transform_service.execute_transform(
        transform_class=transform_class,
        data={
            # Load numpy image directly for processing
            "image": load_original_image_np(transform_class.__name__)
        },
        params=params,
        config=config
    )

    assert "image" in result
    assert isinstance(result["image"], np.ndarray)
    assert result["image"].shape == expected_shape

@pytest.mark.parametrize("transform_class, params", [
    (A.Blur, {"invalid_param": 3}),
    (A.Rotate, {"limit": "invalid_type"}),
])
def test_execute_transform_invalid_params(transform_service, transform_class, params):
    # Create a default config
    config = TransformExecutionConfig(
        targets=["image"],
        bbox_params=None,
        keypoint_params=None,
    )
    with pytest.raises(Exception):
        transform_service.execute_transform(
            transform_class=transform_class,
            data={
                # Load numpy image directly for processing
                "image": load_original_image_np(transform_class.__name__)
            },
            params=params,
            config=config
        )

def test_horizontal_flip_all_targets(transform_service):
    # Define config for HorizontalFlip with all targets
    config = TransformExecutionConfig(
        targets=["image", "mask", "bboxes", "keypoints"],
        bbox_params={"format": "pascal_voc", "label_fields": ["class_labels"], "minArea": 0, "minVisibility": 0},
        keypoint_params={"format": "xy", "label_fields": ["keypoint_labels"], "removeInvisible": False}
    )
    result = transform_service.execute_transform(
        transform_class=A.HorizontalFlip,
        data={
            # Load numpy image directly for processing
            "image": load_original_image_np("HorizontalFlip")
        },
        params={},
        config=config
    )

    # Check that all expected keys are present
    # Adjust assertion based on actual bbox_params used ('class_labels')
    # Also check keypoint label field
    expected_keys = {'image', 'mask', 'bboxes', 'keypoints', 'class_labels', 'keypoint_labels'}
    assert set(result.keys()) >= expected_keys

    # Basic shape checks
    assert isinstance(result['image'], np.ndarray)
    assert result['image'].shape == (733, 484, 3)

    assert isinstance(result['mask'], np.ndarray)
    assert result['mask'].shape == (733, 484)

    assert isinstance(result['bboxes'], list)
    assert len(result['bboxes']) > 0
    for bbox in result['bboxes']:
        assert len(bbox) == 4  # x_min, y_min, x_max, y_max

    assert isinstance(result['keypoints'], list)
    assert len(result['keypoints']) > 0
    for keypoint in result['keypoints']:
        assert len(keypoint) == 2  # x, y

    assert isinstance(result['class_labels'], list)
    assert len(result['class_labels']) == len(result['bboxes'])

    # Verify horizontal flip logic for first bbox and keypoint
    width = load_original_image_np("HorizontalFlip").shape[1]
    first_bbox = result['bboxes'][0]
    first_keypoint = result['keypoints'][0]

    # Check that x coordinates are flipped (x_new = width - x)
    assert first_bbox[0] >= 0 and first_bbox[0] < width  # x_min
    assert first_bbox[2] >= 0 and first_bbox[2] < width  # x_max
    assert first_keypoint[0] >= 0 and first_keypoint[0] < width  # x

@pytest.mark.parametrize("params, expected_params", [
    (
        {
            "interpolation": "cv2.INTER_LINEAR",
            "mask_interpolation": "cv2.INTER_NEAREST",
            "border_mode": "cv2.BORDER_CONSTANT"
        },
        {
            "interpolation": cv2.INTER_LINEAR,
            "mask_interpolation": cv2.INTER_NEAREST,
            "border_mode": cv2.BORDER_CONSTANT
        }
    ),
    (
        {
            "interpolation": "INTER_LINEAR",  # without cv2 prefix
            "scale": 1.0
        },
        {
            "interpolation": cv2.INTER_LINEAR,
            "scale": 1.0
        }
    ),
    (None, {}),  # Test None params
    ({}, {}),  # Test empty params
])
def test_update_transform_params_cv2_constants(transform_service, params, expected_params):
    result = transform_service._update_transform_params(
        params=params,
        transform_class=A.Resize,
    )
    assert result == expected_params

@pytest.mark.parametrize("invalid_params", [
    {"interpolation": "INVALID_INTERPOLATION"},
    {"border_mode": "INVALID_BORDER_MODE"},
    {"mask_interpolation": "cv2.INVALID_INTERPOLATION"},
])
def test_update_transform_params_invalid_cv2_constants(transform_service, invalid_params):
    with pytest.raises(ValueError):
        transform_service._update_transform_params(
            params=invalid_params,
            transform_class=A.Resize,
        )

def test_update_transform_params_with_reference_images(transform_service):
    """Test reference image processing with real images."""
    # Get real images from image service
    # Load and decode reference image
    reference_data_list = image_service.get_reference_image_data("FDA")
    assert reference_data_list
    reference_image = image_service.decode_base64_to_image(reference_data_list[0].image)

    # Test with base64 encoded image
    encoded_reference = image_service.encode_image_to_base64(reference_image)
    params_with_encoded = {
        "reference_images": [encoded_reference],
        "beta_limit": 0.5
    }

    result = transform_service._update_transform_params(
        params=params_with_encoded,
        transform_class=A.FDA,
    )

    assert "reference_images" in result
    assert len(result["reference_images"]) == 1
    # Comment out assertion checking for ndarray, as _update_transform_params no longer decodes
    # assert isinstance(result["reference_images"][0], np.ndarray)
    # assert result["reference_images"][0].shape == reference_image.shape
    # Assert that metadata_key was added
    assert result.get("metadata_key") == "metadata"
    # Assert read_fn is NOT added anymore
    assert "read_fn" not in result

    # Test with no reference images for mixing transform
    result = transform_service._update_transform_params(
        params=None,
        transform_class=A.FDA,
    )

    # Check that metadata_key is added, but reference_images/read_fn are not
    assert result.get("metadata_key") == "metadata"
    assert "reference_images" not in result
    assert "read_fn" not in result

def test_execute_transform_with_cv2_constants(transform_service):
    """Test full transform execution with CV2 constants."""
    # Load numpy image directly for processing
    original_image = load_original_image_np("Resize")

    # Define config for Affine
    config = TransformExecutionConfig(
        targets=["image"],
        bbox_params=None,
        keypoint_params=None,
    )

    result = transform_service.execute_transform(
        transform_class=A.Affine,
        data = {
            "image": original_image
        },
        params={
            "scale": 1.0,
            "translate_percent": 0.0,
            "rotate": 45,
            "shear": 0.0,
            "interpolation": "cv2.INTER_LINEAR",
            "border_mode": "cv2.BORDER_CONSTANT"
        },
        config=config
    )

    assert isinstance(result["image"], np.ndarray)
    assert result["image"].shape == (733, 484, 3)

def test_execute_transform_with_reference_images(transform_service):
    """Test full transform execution with reference images."""
    original_image = load_original_image_np("FDA")
    # Load and decode reference image
    reference_data_list = image_service.get_reference_image_data("FDA")
    assert reference_data_list
    reference_image = image_service.decode_base64_to_image(reference_data_list[0].image)

    # Define config for FDA
    config = TransformExecutionConfig(
        targets=["image"],
        bbox_params=None,
        keypoint_params=None,
    )

    # Test with encoded reference image
    encoded_reference = image_service.encode_image_to_base64(reference_image)
    result = transform_service.execute_transform(
        transform_class=A.FDA,
        data={
            "image": original_image
        },
        params={
            "reference_images": [encoded_reference],
            "beta_limit": 0.5
        },
        config=config
    )

    assert isinstance(result["image"], np.ndarray)
    assert result["image"].shape == original_image.shape

    # Test with numpy array reference image
    result = transform_service.execute_transform(
        transform_class=A.FDA,
        data = {
            "image": original_image
        },
        params={
            "reference_images": [reference_image],
            "beta_limit": 0.5
        },
        config=config
    )

    assert isinstance(result["image"], np.ndarray)
    assert result["image"].shape == original_image.shape

    # Test with default reference image
    result = transform_service.execute_transform(
        transform_class=A.FDA,
        data = {
            "image": original_image
        },
        params={
            "beta_limit": 0.5
        },
        config=config
    )

    assert isinstance(result["image"], np.ndarray)
    assert result["image"].shape == original_image.shape
