"""Mask2Former instance-segmentation adapter.

Backbone: `transformers.AutoModelForUniversalSegmentation` on
`CondadosAI/mask2former_swin_tiny_coco_instance` (mirror of
`facebook/mask2former-swin-tiny-coco-instance`).

    predict flow
    ────────────
    inputs
       │  io.load_images(strict_rgb=...)
       ▼
    [PIL RGB] x N
       │  chunk into batches of `batch_size`
       ▼
    each batch
       │  processor(images=batch, return_tensors="pt") → .to(device)
       ▼
    with torch.inference_mode():
        outputs = model(**batch_inputs)
       │
       ▼
    processor.post_process_instance_segmentation(
        outputs, target_sizes=[(h, w), ...], threshold=0.5
    ) → per-image {"segmentation": Tensor[H,W] int, "segments_info": [...]}
       │
       │  (transformers returns a per-pixel instance-id map; acaua's
       │   SegmentationResult shape wants per-instance masks + boxes)
       │  _unpack_instance_seg() splits the map into N binary masks and
       │  _masks_to_boxes() computes tight xyxy bboxes.
       ▼
    SegmentationResult per image (boxes, labels, scores, masks)
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import torch
from transformers import AutoImageProcessor, AutoModelForUniversalSegmentation

from acaua.io import ImageInput, load_images
from acaua.model import Model
from acaua.results import Result, SegmentationResult

_DEFAULT_THRESHOLD = 0.5


class Mask2FormerAdapter(Model):
    """Mask2Former instance-segmentation adapter."""

    def __init__(
        self,
        *,
        repo_id: str,
        processor: Any,
        model: Any,
        device: torch.device,
    ) -> None:
        self.repo_id = repo_id
        self._processor = processor
        self._model = model
        self.device = device

    @classmethod
    def _load(
        cls,
        *,
        name: str,
        revision: str,
        device: torch.device,
    ) -> Mask2FormerAdapter:
        processor = AutoImageProcessor.from_pretrained(name, revision=revision)  # type: ignore[no-untyped-call]
        model = AutoModelForUniversalSegmentation.from_pretrained(name, revision=revision)
        model.eval()
        model.to(device)
        return cls(repo_id=name, processor=processor, model=model, device=device)

    def predict(
        self,
        inputs: ImageInput | Sequence[ImageInput],
        *,
        batch_size: int = 8,
        strict_rgb: bool = True,
    ) -> Result | list[Result]:
        pil_images, was_sequence = load_images(inputs, strict_rgb=strict_rgb)

        results: list[Result] = []
        for start in range(0, len(pil_images), batch_size):
            batch = pil_images[start : start + batch_size]
            target_sizes = [(img.height, img.width) for img in batch]
            batch_inputs = self._processor(images=batch, return_tensors="pt")
            batch_inputs = {k: v.to(self.device) for k, v in batch_inputs.items()}
            with torch.inference_mode():
                outputs = self._model(**batch_inputs)
            per_image = self._processor.post_process_instance_segmentation(
                outputs,
                target_sizes=target_sizes,
                threshold=_DEFAULT_THRESHOLD,
            )
            for per_image_result in per_image:
                results.append(_unpack_instance_seg(per_image_result, self.device))

        return results if was_sequence else results[0]


# ---------------------------------------------------------------------------
# post-processing helpers
# ---------------------------------------------------------------------------


def _unpack_instance_seg(result: dict[str, Any], device: torch.device) -> SegmentationResult:
    """Convert transformers instance-seg output to a `SegmentationResult`.

    transformers returns one dict per image:
        segmentation:  Tensor[H, W] int — each pixel's instance id (bg = -1)
        segments_info: list of {id, label_id, score}

    acaua's SegmentationResult wants per-instance tensors. We split the
    segmentation map into N binary masks and compute tight xyxy bboxes
    from each mask's extents.
    """
    segmentation = result["segmentation"].to(device)
    segments_info = result["segments_info"]
    h, w = segmentation.shape[-2], segmentation.shape[-1]

    if len(segments_info) == 0:
        return SegmentationResult(
            boxes=torch.zeros(0, 4, dtype=torch.float32, device=device),
            labels=torch.zeros(0, dtype=torch.int64, device=device),
            scores=torch.zeros(0, dtype=torch.float32, device=device),
            masks=torch.zeros(0, h, w, dtype=torch.bool, device=device),
        )

    masks = torch.stack([segmentation == seg["id"] for seg in segments_info], dim=0)
    labels = torch.tensor(
        [int(seg["label_id"]) for seg in segments_info],
        dtype=torch.int64,
        device=device,
    )
    scores = torch.tensor(
        [float(seg["score"]) for seg in segments_info],
        dtype=torch.float32,
        device=device,
    )
    boxes = _masks_to_boxes(masks)
    return SegmentationResult(boxes=boxes, labels=labels, scores=scores, masks=masks)


def _masks_to_boxes(masks: torch.Tensor) -> torch.Tensor:
    """Tight xyxy bboxes from per-instance binary masks.

    Input:  `[N, H, W]` bool tensor (one mask per instance).
    Output: `[N, 4]` float32 on the same device. Empty masks → zero row.
    """
    n = masks.shape[0]
    boxes = torch.zeros(n, 4, dtype=torch.float32, device=masks.device)
    for i in range(n):
        ys, xs = masks[i].nonzero(as_tuple=True)
        if ys.numel() == 0:
            continue
        boxes[i, 0] = xs.min().item()
        boxes[i, 1] = ys.min().item()
        boxes[i, 2] = xs.max().item() + 1
        boxes[i, 3] = ys.max().item() + 1
    return boxes
