import numpy as np
import pytest
from fastapi.testclient import TestClient

from src.main import app
from src.services.image_service import ImageService
from src.config import settings

# Define client and service instances globally for this test module
client = TestClient(app)
image_service = ImageService()

# Helper to get encoded original image data string
def get_encoded_original_image(transform_name: str) -> str:
    return image_service.get_original_image_data(transform_name).image

# 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.mark.parametrize("transform_data", [
    {
        "name": "Blur",
        "params": {"blur_limit": 7},
        "targets": ["image"]
    },
    {
        "name": "Affine",
        "params": {"scale": 1.2, "translate_percent": {"x": 0.1, "y": 0.1}, "rotate": 45},
        "targets": ["image", "mask", "bboxes", "keypoints"]
    },
    {
        "name": "FDA",
        "params": {
            "reference_images": None,  # Will be filled with encoded test image
            "beta_limit": 0.4
        },
        "targets": ["image"]
    },
    {
        "name": "Mosaic",
        "params": {
        },
        "targets": ["image", "mask", "bboxes", "keypoints"]
    }
])
def test_execute_transform_success(transform_data):
    """Test successful transform execution"""
    # Use helper to load numpy array for processing
    test_image_np = load_original_image_np(transform_data["name"])
    encoded_image = image_service.encode_image_to_base64(test_image_np)

    # Update FDA params with encoded reference image if needed
    if transform_data["name"] == "FDA":
        transform_data["params"]["reference_images"] = [image_service.encode_image_to_base64(test_image_np)]

    request_data = {
        "transform": transform_data["name"],
        "data": {"image": encoded_image},
        "params": transform_data["params"],
        "targets": transform_data["targets"]
    }

    if transform_data["name"] in {"Affine", "Mosaic"}:
        request_data["bbox_params"] = {
            "format": "pascal_voc",
        }
        request_data["keypoint_params"] = {
            "format": "xy",
        }

    response = client.post("/execute-transform", json=request_data)

    assert response.status_code == 200, response.json()
    response_data = response.json()
    assert "image" in response_data



@pytest.mark.parametrize("transform_data", [
    {
        "name": "ColorJitter",
        "params": {"brightness": -1.0, "contrast": 1.0, "saturation": 1.0, "hue": 0.1},
        "targets": ["image"]
    },
    {
        "name": "Affine",
        "params": {"scale": "a"},
        "targets": ["image", "mask", "bboxes", "keypoints"]
    },
    {
        "name": "FDA",
        "params": {"reference_image": None},
        "targets": ["image", "mask", "bboxes", "keypoints"]
    }
])
def test_execute_transform_errors(transform_data):
    """Test transform execution errors"""
    request_data = {
        "transform": transform_data["name"],
        # Use helper to get encoded image string for request
        "data": {"image": get_encoded_original_image(transform_data["name"])},
        "params": transform_data["params"],
        "targets": transform_data["targets"]
    }

    # Add bbox_params and keypoint_params if needed for Affine transform
    if transform_data["name"] == "Affine" and "bboxes" in transform_data["targets"]:
        request_data["bbox_params"] = {
            "format": "pascal_voc",
        }

    if transform_data["name"] == "Affine" and "keypoints" in transform_data["targets"]:
        request_data["keypoint_params"] = {
            "format": "xy",
        }

    response = client.post("/execute-transform", json=request_data)

    # Check that we get an error response (either 400 or 500)
    assert response.status_code in (400, 500), f"Expected error status code, got {response.status_code}"

def test_execute_transform_no_image():
    """Test transform execution with no image (should use default)"""
    request_data = {
        "transform": "Blur",
        "params": {"blur_limit": 7},
        "targets": ["image"],
        "data": {}
    }

    response = client.post("/execute-transform", json=request_data)
    assert response.status_code == 200
    response_data = response.json()
    assert "image" in response_data


def test_execute_transform_invalid_base64():
    """Test execution with invalid base64 image"""
    request_data = {
        "transform": "Blur",
        "data": {"image": "invalid_base64_string"},
        "params": {"blur_limit": 7},
        "targets": ["image"]
    }

    response = client.post("/execute-transform", json=request_data)
    assert response.status_code == 500, response.json()
    assert "Invalid image data" in response.json()["detail"]

@pytest.mark.parametrize("transform", ["FDA"])
def test_execute_transform_default_reference_image(transform):
    """Test transform execution with default reference image"""
    request_data = {
        "transform": transform,
        # Use helper to get encoded image string for request
        "data": {"image": get_encoded_original_image(transform)},
        "params": {"reference_images": None},
        "targets": ["image"]
    }

    response = client.post("/execute-transform", json=request_data)
    response_data = response.json()
    assert response.status_code == 200, response_data
    assert "image" in response_data

@pytest.mark.parametrize("transform_config", [
    {
        "name": "Blur",
        "params": {"blur_limit": 7},
        "needs_reference": False,
        "targets": ["image"]
    },
    {
        "name": "Affine",
        "params": {"scale": 1.2, "translate_percent": {"x": 0.1, "y": 0.1}, "rotate": 45},
        "needs_reference": False,
        "targets": ["image", "mask", "bboxes", "keypoints"]
    },
    {
        "name": "FDA",
        "params": {"beta_limit": 0.4},
        "needs_reference": True,
        "targets": ["image"]
    }
])
@pytest.mark.parametrize("provide_image", [True, False])
@pytest.mark.parametrize("provide_reference", [True, False])
def test_transform_combinations(transform_config, provide_image, provide_reference):
    """Test all combinations of transforms with different image configurations"""
    # Skip non-relevant reference image combinations
    if not transform_config["needs_reference"] and provide_reference:
        pytest.skip("This transform doesn't use reference images")

    # Prepare request data
    request_data = {
        "transform": transform_config["name"],
        "params": transform_config["params"].copy(),
        "targets": transform_config["targets"],
        "data": {}
    }

    # Add main image if specified
    if provide_image:
        # Use helper to get encoded image string for request
        request_data["data"] = {"image": get_encoded_original_image(transform_config["name"])}

    # Add reference image if needed and specified
    if transform_config["needs_reference"]:
        if provide_reference:
            # Use helper to get encoded image string for request
            request_data["params"]["reference_images"] = [get_encoded_original_image(transform_config["name"])]
        else:
            request_data["params"]["reference_images"] = None


    if transform_config["name"] == "Affine":
        request_data["bbox_params"] = {
            "format": "pascal_voc",
        }
        request_data["keypoint_params"] = {
            "format": "xy",
            "remove_invisible": True
        }

    # Make request
    response = client.post("/execute-transform", json=request_data)

    # Assertions
    assert response.status_code == 200
    response_data = response.json()

    # Verify response structure
    assert "image" in response_data

    # Verify image can be decoded
    result_image = image_service.decode_base64_to_image(response_data["image"])
    assert len(result_image.shape) == 3  # Should be RGB
    assert result_image.shape[2] == 3    # Should have 3 channels

def test_invalid_transform():
    """Test execution with invalid transform name"""
    request_data = {
        "transform": "NonExistentTransform",
        # Use helper to get encoded image string for request
        "data": {"image": get_encoded_original_image("Blur")}, # Use a valid transform name here
        "params": {},
        "targets": ["image"]
    }

    response = client.post("/execute-transform", json=request_data)
    assert response.status_code == 500
    assert "Transform not found" in response.json()["detail"]


def test_affine_transform_visualizations():
    """Test that Affine transform returns all expected visualization elements"""
    # Use helper to get encoded image string for request
    encoded_image = get_encoded_original_image("Affine")

    request_data = {
        "transform": "Affine",
        "data": {"image": encoded_image},
        "targets": ["image", "mask", "bboxes", "keypoints"],
        "params": {
            "scale": 1.2,
            "translate_percent": {"x": 0.1, "y": 0.1},
            "rotate": 45
        },
        "bbox_params": {
            "format": "pascal_voc",
            "label_fields": ["labels"],
            "labels": ["test"]
        },
        "keypoint_params": {
            "format": "xy",
        }
    }

    response = client.post("/execute-transform", json=request_data)
    assert response.status_code == 200

    response_data = response.json()

    # Check all expected fields are present
    expected_fields = [
        "image",
        "mask",
        "bboxes",
        "keypoints"
    ]

    for field in expected_fields:
        assert field in response_data, f"Missing {field} in response"

    # Verify each image field can be decoded
    for field in ["image", "mask"]:
        decoded_image = image_service.decode_base64_to_image(response_data[field])
        assert len(decoded_image.shape) == 3  # Should be RGB
        assert decoded_image.shape[2] == 3    # Should have 3 channels


@pytest.mark.parametrize("test_case", [
    {
        "targets": ["image", "bboxes"],
        "include_bbox_params": False,
        "include_keypoint_params": False,
        "expected_status_code": 400,
        "expected_error": "bbox_params must be provided"
    },
    {
        "targets": ["image", "keypoints"],
        "include_bbox_params": False,
        "include_keypoint_params": False,
        "expected_status_code": 400,
        "expected_error": "keypoint_params must be provided"
    },
    {
        "targets": ["image", "bboxes", "keypoints"],
        "include_bbox_params": True,
        "include_keypoint_params": False,
        "expected_status_code": 400,
        "expected_error": "keypoint_params must be provided"
    },
    {
        "targets": ["image", "bboxes", "keypoints"],
        "include_bbox_params": False,
        "include_keypoint_params": True,
        "expected_status_code": 400,
        "expected_error": "bbox_params must be provided"
    },
    {
        "targets": ["image", "bboxes", "keypoints"],
        "include_bbox_params": True,
        "include_keypoint_params": True,
        "expected_status_code": 200,
        "expected_error": None
    },
    {
        "targets": ["image"],
        "include_bbox_params": False,
        "include_keypoint_params": False,
        "expected_status_code": 200,
        "expected_error": None
    }
])
def test_required_params_validation(test_case):
    """Test that API validates required parameters for bboxes and keypoints"""
    # Prepare request data
    request_data = {
        "transform": "Affine",
        "data": {"image": get_encoded_original_image("Affine")},
        "params": {"scale": 1.2, "translate_percent": {"x": 0.1, "y": 0.1}, "rotate": 45},
        "targets": test_case["targets"]
    }

    # Add bbox_params if needed
    if test_case["include_bbox_params"]:
        request_data["bbox_params"] = {
            "format": "pascal_voc",
        }

    # Add keypoint_params if needed
    if test_case["include_keypoint_params"]:
        request_data["keypoint_params"] = {
            "format": "xy",
        }

    # Make request
    response = client.post("/execute-transform", json=request_data)

    # Check status code
    assert response.status_code == test_case["expected_status_code"]

    # Check error message if expected
    if test_case["expected_error"]:
        assert test_case["expected_error"] in response.json()["detail"]


def test_execute_transform_with_custom_image():
    """Test that the API processes the provided image rather than using the default image."""
    # Create a custom image (a simple red square)
    np.random.seed(137)
    custom_image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)

    # Test the encoding/decoding process directly
    encoded_image = image_service.encode_image_to_base64(custom_image)
    decoded_image = image_service.decode_base64_to_image(encoded_image)

    encode_decode_diff = np.mean(np.abs(decoded_image.astype(float) - custom_image.astype(float)))

    # Create request with NoOp transform (which should return the image unchanged)
    request_data = {
        "transform": "NoOp",
        "data": {"image": encoded_image},
        "params": {},
        "targets": ["image"]
    }

    # Execute the transform
    response = client.post("/execute-transform", json=request_data)

    # Verify response
    assert response.status_code == 200
    response_data = response.json()

    # Decode the returned image
    returned_image = image_service.decode_base64_to_image(response_data["image"])

    # Check if the returned image has the same dimensions as our custom image
    assert returned_image.shape == custom_image.shape, "Returned image should have the same shape as our custom image"

    # Get the default image data and decode it
    default_image_data = image_service.get_original_image_data("NoOp")
    default_image = image_service.decode_base64_to_image(default_image_data.image)

    # Verify the shapes are different (this proves we're not getting the default image)
    assert returned_image.shape != default_image.shape, "Returned image should have different shape than default image"

    api_diff = np.mean(np.abs(returned_image.astype(float) - custom_image.astype(float)))

    # Calculate the difference between the returned image and the default image
    # This should be large since they're completely different images
    default_diff = np.mean(np.abs(
        returned_image[:min(returned_image.shape[0], default_image.shape[0]),
                      :min(returned_image.shape[1], default_image.shape[1]),
                      :].astype(float) -
        default_image[:min(returned_image.shape[0], default_image.shape[0]),
                     :min(returned_image.shape[1], default_image.shape[1]),
                     :].astype(float)
    ))

    # The key assertions:

    # 1. The API difference should be similar to the encode/decode difference
    # This confirms that the API is just encoding/decoding our image, not modifying it significantly
    assert abs(api_diff - encode_decode_diff) < 10.0, \
        f"API difference ({api_diff}) should be similar to encode/decode difference ({encode_decode_diff})"

    # 2. The correlation between original and returned images should be high
    # This confirms that the images are visually similar despite JPEG artifacts
    correlation = np.corrcoef(custom_image.reshape(-1), returned_image.reshape(-1))[0, 1]
    assert correlation > 0.5, f"Correlation too low: {correlation}"
