"""SAM mask proposal generation via bounding box prompts."""

from typing import Any, Sequence

import numpy as np
from PIL import Image
from transformers import Sam2Model, Sam2Processor
from utils import get_device


def generate_proposals(
    model: Sam2Model,
    processor: Sam2Processor,
    image: Image.Image,
    bboxes: list[list[float]],
) -> list[dict[str, Any]]:
    """Segment regions in image using SAM2 with bounding box prompts.

    Args:
        model: Sam2Model instance.
        processor: Sam2Processor instance.
        image: PIL Image to segment.
        bboxes: Bounding boxes as [[x1, y1, x2, y2], ...].

    Returns:
        List of mask dictionaries with keys:
            - segment: Binary mask (numpy array)
            - area: Mask area in pixels
            - bbox: Bounding box [x, y, width, height]
            - predicted_iou: Model's confidence in the mask
            - stability_score: Stability score for the mask
    """
    if not bboxes:
        return []

    device = get_device()
    image_rgb = image.convert("RGB")
    input_boxes = [bboxes]
    inputs = processor(
        images=image_rgb,
        input_boxes=input_boxes,
        return_tensors="pt",
    ).to(device)

    outputs = model(**inputs, multimask_output=False)
    masks = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"],
    )[0]

    return _masks_to_proposals(masks)


def generate_proposals_batch(
    model: Sam2Model,
    processor: Sam2Processor,
    images: Sequence[Image.Image],
    bboxes_per_image: list[list[list[float]]],
) -> list[list[dict[str, Any]]]:
    """Segment a batch of images using SAM2 with bounding box prompts.

    Args:
        model: Sam2Model instance.
        processor: Sam2Processor instance.
        images: Sequence of PIL Images to segment.
        bboxes_per_image: Bounding boxes per image, outer list matches images length.

    Returns:
        List of lists of mask dictionaries, one inner list per image.
    """
    image_list = list(images)
    if not image_list:
        return []

    device = get_device()
    image_rgb_list = [img.convert("RGB") for img in image_list]

    inputs = processor(
        images=image_rgb_list,
        input_boxes=bboxes_per_image,
        return_tensors="pt",
    ).to(device)

    outputs = model(**inputs, multimask_output=False)
    all_masks = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"],
    )

    return [_masks_to_proposals(image_masks) for image_masks in all_masks]


def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]:
    """Convert model output masks to list of mask dicts."""
    mask_array = _to_numpy_mask_array(masks)
    if mask_array is None:
        return []

    # Ensure 3D: [num_masks, H, W]
    if mask_array.ndim == 2:
        mask_array = np.expand_dims(mask_array, axis=0)

    if mask_array.ndim != 3:
        return []

    # Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W]
    if mask_array.ndim == 3 and mask_array.shape[0] == 1:
        mask_array = mask_array[0]

    proposals: list[dict[str, Any]] = []
    for single_mask in mask_array:
        mask_dict = _build_mask_dict(single_mask)
        if mask_dict is not None:
            proposals.append(mask_dict)

    return proposals


def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
    """Convert mask-like object to numpy array."""
    if mask_like is None:
        return None
    if isinstance(mask_like, np.ndarray):
        return mask_like

    import torch

    if isinstance(mask_like, torch.Tensor):
        return mask_like.detach().cpu().numpy()

    return None


def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
    """Build a mask dictionary from a 2D boolean numpy array."""
    if mask_array.ndim != 2:
        return None
    segment = mask_array.astype(bool)
    area = int(segment.sum())
    if area <= 0:
        return None

    ys, xs = np.where(segment)
    min_y, max_y = int(ys.min()), int(ys.max())
    min_x, max_x = int(xs.min()), int(xs.max())
    bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]

    return {
        "segment": segment,
        "area": area,
        "bbox": bbox,
        "predicted_iou": None,
        "stability_score": None,
    }