#!/usr/bin/env python3
"""
Test script for YOLO preprocessing of 16-bit TIFF images.
"""

import numpy as np
import tifffile
from pathlib import Path
import tempfile
import sys
import os

# Add parent directory to path to import modules
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.model.yolo_wrapper import YOLOWrapper
from src.utils.image import Image
from PIL import Image as PILImage


def create_test_16bit_tiff(output_path: str) -> str:
    """Create a test 16-bit grayscale TIFF file.

    Args:
        output_path: Path where to save the test TIFF

    Returns:
        Path to the created TIFF file
    """
    # Create a 16-bit grayscale test image (200x200)
    # With values ranging from 0 to 65535 (full 16-bit range)
    height, width = 200, 200

    # Create a gradient pattern
    test_data = np.zeros((height, width), dtype=np.uint16)
    for i in range(height):
        for j in range(width):
            # Create a diagonal gradient
            test_data[i, j] = int((i + j) / (height + width - 2) * 65535)

    # Save as TIFF
    tifffile.imwrite(output_path, test_data)
    print(f"Created test 16-bit TIFF: {output_path}")
    print(f"  Shape: {test_data.shape}")
    print(f"  Dtype: {test_data.dtype}")
    print(f"  Min value: {test_data.min()}")
    print(f"  Max value: {test_data.max()}")

    return output_path


def test_yolo_preprocessing():
    """Test YOLO preprocessing of 16-bit TIFF images."""
    print("\n=== Testing YOLO Preprocessing of 16-bit TIFF ===")

    # Create temporary test file
    with tempfile.NamedTemporaryFile(suffix=".tif", delete=False) as tmp:
        test_path = tmp.name

    try:
        # Create test image
        create_test_16bit_tiff(test_path)

        # Create YOLOWrapper instance (no actual model loading needed for this test)
        print("\nTesting YOLOWrapper._prepare_source()...")
        wrapper = YOLOWrapper()

        # Call _prepare_source to preprocess the image
        prepared_path, cleanup_path = wrapper._prepare_source(test_path)

        print(f"\nPreprocessing complete:")
        print(f"  Original path: {test_path}")
        print(f"  Prepared path: {prepared_path}")
        print(f"  Cleanup path: {cleanup_path}")

        # Verify the prepared image exists
        assert os.path.exists(prepared_path), "Prepared image should exist"

        # Load the prepared image and verify it's uint8 RGB
        prepared_img = PILImage.open(prepared_path)
        print(f"\nPrepared image properties:")
        print(f"  Mode: {prepared_img.mode}")
        print(f"  Size: {prepared_img.size}")
        print(f"  Format: {prepared_img.format}")

        # Convert to numpy to check values
        img_array = np.array(prepared_img)
        print(f"  Shape: {img_array.shape}")
        print(f"  Dtype: {img_array.dtype}")
        print(f"  Min value: {img_array.min()}")
        print(f"  Max value: {img_array.max()}")
        print(f"  Mean value: {img_array.mean():.2f}")

        # Verify it's RGB uint8
        assert prepared_img.mode == "RGB", "Prepared image should be RGB"
        assert img_array.dtype == np.uint8, "Prepared image should be uint8"
        assert img_array.shape[2] == 3, "Prepared image should have 3 channels"
        assert (
            0 <= img_array.min() <= img_array.max() <= 255
        ), "Values should be in [0, 255]"

        # Cleanup prepared file if needed
        if cleanup_path and os.path.exists(cleanup_path):
            os.remove(cleanup_path)
            print(f"\nCleaned up prepared image: {cleanup_path}")

        print("\n✓ All YOLO preprocessing tests passed!")
        return True

    except Exception as e:
        print(f"\n✗ Test failed with error: {e}")
        import traceback

        traceback.print_exc()
        return False

    finally:
        # Cleanup
        if os.path.exists(test_path):
            os.remove(test_path)
            print(f"Cleaned up test file: {test_path}")


if __name__ == "__main__":
    success = test_yolo_preprocessing()
    sys.exit(0 if success else 1)