"""Image loading and processing utilities."""

import os
import numpy as np
from tifffile import imread, TiffFile
import logging

logger = logging.getLogger(__name__)

# Auto-downsample factor from environment (0=disabled, 2=2x, 4=4x, etc.)
AUTO_DOWNSAMPLE = int(os.getenv("AIMINO_AUTO_DOWNSAMPLE", "0"))


def _to_2d_gray_safe(arr):
    """Convert array to 2D grayscale, handling various input shapes."""
    a = np.squeeze(arr)
    if a.ndim == 2:
        return a
    if a.ndim == 3 and a.shape[-1] in (3, 4):
        rgb = a[..., :3].astype(float)
        return (
            0.2989 * rgb[..., 0] + 0.587 * rgb[..., 1] + 0.114 * rgb[..., 2]
        ).astype(a.dtype)
    return _to_2d_gray_safe(a[0])


def _apply_downsample(img: np.ndarray, factor: int) -> np.ndarray:
    """Downsample image by given factor."""
    if factor <= 1:
        return img
    logger.info(f"[image] auto-downsampling by {factor}x: {img.shape} -> {img.shape[0]//factor}x{img.shape[1]//factor}")
    return img[::factor, ::factor]


def load_image_for_mask(path: str) -> np.ndarray:
    """Load image from TIFF file for mask processing."""
    logger.info(f"[image] loading image for mask from {path}")
    with TiffFile(path) as tf:
        arr = tf.series[0].asarray()
    img = _to_2d_gray_safe(arr)

    # Auto-downsample if enabled and image is large
    if AUTO_DOWNSAMPLE > 1:
        img = _apply_downsample(img, AUTO_DOWNSAMPLE)
    elif max(img.shape) > 16384:
        # Auto-calculate downsample factor to fit within 16k
        factor = (max(img.shape) // 16384) + 1
        logger.warning(f"[image] Image {img.shape} exceeds 16k, auto-downsampling by {factor}x")
        img = _apply_downsample(img, factor)

    return img

