"""
Tests for image and mask loading module.
"""

import pytest
import numpy as np
from pathlib import Path
import tempfile

from medtda.loaders import (
    load_image,
    load_mask,
    validate_image,
    validate_mask,
    validate_compatibility,
    save_image,
)


class TestValidateImage:
    """Tests for validate_image function."""
    
    def test_valid_2d_image(self):
        """Test validation of valid 2D image."""
        image = np.random.rand(50, 50)
        validate_image(image)  # Should not raise
    
    def test_valid_3d_image(self):
        """Test validation of valid 3D image."""
        image = np.random.rand(30, 30, 30)
        validate_image(image)  # Should not raise
    
    def test_valid_4d_image(self):
        """Test validation of valid 4D image."""
        image = np.random.rand(20, 20, 20, 10)
        validate_image(image)  # Should not raise
    
    def test_invalid_type(self):
        """Test that non-array raises TypeError."""
        with pytest.raises(TypeError, match="must be a numpy array"):
            validate_image([1, 2, 3])
    
    def test_invalid_dimensions(self):
        """Test that 1D array raises ValueError."""
        image = np.array([1, 2, 3, 4, 5])
        with pytest.raises(ValueError, match="must be 2D, 3D, or 4D"):
            validate_image(image)
    
    def test_empty_image(self):
        """Test that empty array raises ValueError."""
        image = np.array([])
        with pytest.raises(ValueError, match="array is empty"):
            validate_image(image)
    
    def test_nan_values(self):
        """Test that NaN values raise ValueError."""
        image = np.random.rand(20, 20)
        image[10, 10] = np.nan
        with pytest.raises(ValueError, match="NaN or infinite"):
            validate_image(image)
    
    def test_inf_values(self):
        """Test that infinite values raise ValueError."""
        image = np.random.rand(20, 20)
        image[5, 5] = np.inf
        with pytest.raises(ValueError, match="NaN or infinite"):
            validate_image(image)


class TestValidateMask:
    """Tests for validate_mask function."""
    
    def test_valid_binary_mask(self):
        """Test validation of valid binary mask."""
        mask = np.random.randint(0, 2, size=(50, 50), dtype=np.uint8)
        validate_mask(mask)  # Should not raise
    
    def test_invalid_type(self):
        """Test that non-array raises TypeError."""
        with pytest.raises(TypeError, match="must be a numpy array"):
            validate_mask([0, 1, 1, 0])
    
    def test_non_binary_values(self):
        """Test that non-binary values raise ValueError."""
        mask = np.array([[0, 1, 2], [1, 0, 1]], dtype=np.uint8)
        with pytest.raises(ValueError, match="must be binary"):
            validate_mask(mask)


class TestValidateCompatibility:
    """Tests for validate_compatibility function."""
    
    def test_compatible_shapes(self):
        """Test that compatible shapes pass validation."""
        image = np.random.rand(50, 50)
        mask = np.random.randint(0, 2, size=(50, 50), dtype=np.uint8)
        validate_compatibility(image, mask)  # Should not raise
    
    def test_incompatible_shapes(self):
        """Test that incompatible shapes raise ValueError."""
        image = np.random.rand(50, 50)
        mask = np.random.randint(0, 2, size=(40, 40), dtype=np.uint8)
        with pytest.raises(ValueError, match="does not match"):
            validate_compatibility(image, mask)


class TestLoadSave2DImage:
    """Tests for loading and saving 2D images."""
    
    @pytest.mark.skipif(
        not pytest.importorskip("PIL", reason="Pillow not installed"),
        reason="Pillow required for 2D image support"
    )
    def test_save_and_load_png(self):
        """Test saving and loading PNG image."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create test image
            image = np.random.rand(50, 50) * 255
            output_path = Path(tmpdir) / "test_image.png"
            
            # Save image
            save_image(image, output_path)
            
            # Load image back
            loaded_image, metadata = load_image(output_path)
            
            # Check properties
            assert loaded_image.shape == image.shape
            assert metadata['format'] == '.png'
            assert metadata['spacing'] == (1.0, 1.0)


class TestLoadMask:
    """Tests for load_mask function."""
    
    @pytest.mark.skipif(
        not pytest.importorskip("PIL", reason="Pillow not installed"),
        reason="Pillow required for 2D image support"
    )
    def test_load_mask_preserves_labels(self):
        """Test that mask label values are preserved (multi-label support)."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create test mask with various label values
            from PIL import Image
            mask_array = np.array([[0, 1, 2], [3, 2, 1]], dtype=np.uint8)
            mask_img = Image.fromarray(mask_array)
            mask_path = Path(tmpdir) / "test_mask.png"
            mask_img.save(mask_path)
            
            # Load mask
            loaded_mask, metadata = load_mask(mask_path)
            
            # Check that label values are preserved (not binarized)
            unique_values = np.unique(loaded_mask)
            assert np.all(np.isin(unique_values, [0, 1, 2, 3]))
            assert loaded_mask.dtype == np.uint8


class TestFileNotFound:
    """Tests for file not found errors."""
    
    def test_load_image_file_not_found(self):
        """Test that loading non-existent file raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError):
            load_image("nonexistent_file.png")
    
    def test_load_mask_file_not_found(self):
        """Test that loading non-existent mask raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError):
            load_mask("nonexistent_mask.png")


class TestUnsupportedFormat:
    """Tests for unsupported file format errors."""
    
    def test_unsupported_extension(self):
        """Test that unsupported extension raises ValueError."""
        with tempfile.TemporaryDirectory() as tmpdir:
            fake_file = Path(tmpdir) / "test.xyz"
            fake_file.touch()
            
            with pytest.raises(ValueError, match="Unsupported file format"):
                load_image(fake_file)
