"""Mask postprocessing using numpy + Pillow (replaces F.interpolate + threshold)."""

import logging

import numpy as np
from PIL import Image

from sam2_onnx.types import Float32Array, HW

logger = logging.getLogger(__name__)

_MAX_OUTPUT_DIM = 16384


def upscale_masks(
    low_res_masks: Float32Array,
    target_hw: HW,
) -> Float32Array:
    """Bilinear upscale low-resolution mask logits.

    Args:
        low_res_masks: (B, M, H, W) float32 logits.
        target_hw: Target (height, width).

    Returns:
        (B, M, target_h, target_w) float32 logits.
    """
    b, m, _, _ = low_res_masks.shape
    th, tw = target_hw
    if th <= 0 or tw <= 0:
        raise ValueError(f"target_hw dimensions must be positive, got {target_hw}")
    if th > _MAX_OUTPUT_DIM or tw > _MAX_OUTPUT_DIM:
        raise ValueError(
            f"target_hw dimensions exceed maximum {_MAX_OUTPUT_DIM}, got {target_hw}"
        )
    out = np.empty((b, m, th, tw), dtype=np.float32)
    for i in range(b):
        for j in range(m):
            pil_mask = Image.fromarray(low_res_masks[i, j])
            resized = pil_mask.resize((tw, th), Image.Resampling.BILINEAR)
            out[i, j] = np.array(resized, dtype=np.float32)
    return out


def postprocess_masks(
    low_res_masks: Float32Array,
    orig_hw: HW,
    threshold: float = 0.0,
    return_logits: bool = False,
) -> Float32Array:
    """Upscale and threshold masks.

    Args:
        low_res_masks: (B, M, H, W) float32 logits from decoder.
        orig_hw: Original image (height, width).
        threshold: Mask threshold (default 0.0).
        return_logits: If True, return logits without thresholding.

    Returns:
        (B, M, orig_h, orig_w) float32 masks (0.0/1.0) or logits.
    """
    masks = upscale_masks(low_res_masks, orig_hw)
    if not return_logits:
        masks = (masks > threshold).astype(np.float32)
    logger.debug("postprocessed: %s -> %s", low_res_masks.shape, masks.shape)
    return masks
