"""Unit tests for `Mask2FormerAdapter` with mocked transformers calls.

Network-free. Covers:
  - `_load` plumbing (revision + eval + .to).
  - `predict` return shape (single vs list) and batching.
  - Forward pass runs under `torch.inference_mode()`.
  - `_unpack_instance_seg` converts transformers per-pixel segmentation maps
    into acaua's per-instance (box, label, score, mask) shape.
  - `_masks_to_boxes` produces tight xyxy bboxes, handles empty masks.
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
import torch

from acaua import SegmentationResult
from acaua.adapters.mask2former import (
    Mask2FormerAdapter,
    _masks_to_boxes,
    _unpack_instance_seg,
)


def _make_processor_mock(n_instances_per_image: int = 2, h: int = 32, w: int = 32) -> MagicMock:
    processor = MagicMock()

    def call(images, return_tensors):  # type: ignore[no-untyped-def]
        b = len(images)
        return {"pixel_values": torch.zeros(b, 3, h, w)}

    processor.side_effect = call

    def post_process(outputs, target_sizes, threshold):  # type: ignore[no-untyped-def]
        results = []
        for _ in target_sizes:
            seg = torch.full((h, w), -1, dtype=torch.int64)
            # Stamp a few distinct instance ids into the map.
            for inst_id in range(n_instances_per_image):
                y0 = inst_id * (h // (n_instances_per_image + 1))
                y1 = y0 + 4
                seg[y0:y1, 4:12] = inst_id
            segments_info = [
                {"id": inst_id, "label_id": inst_id + 1, "score": 0.9}
                for inst_id in range(n_instances_per_image)
            ]
            results.append({"segmentation": seg, "segments_info": segments_info})
        return results

    processor.post_process_instance_segmentation = MagicMock(side_effect=post_process)
    return processor


def _make_model_mock() -> MagicMock:
    model = MagicMock()
    model.eval = MagicMock(return_value=model)
    model.to = MagicMock(return_value=model)
    model.grad_enabled_during_forward: list[bool] = []

    def forward(**kwargs):  # type: ignore[no-untyped-def]
        model.grad_enabled_during_forward.append(torch.is_grad_enabled())
        return MagicMock()

    model.side_effect = forward
    return model


@pytest.fixture
def mocked_adapter() -> Mask2FormerAdapter:
    processor = _make_processor_mock()
    model = _make_model_mock()
    with (
        patch(
            "acaua.adapters.mask2former.AutoImageProcessor.from_pretrained",
            return_value=processor,
        ),
        patch(
            "acaua.adapters.mask2former.AutoModelForUniversalSegmentation.from_pretrained",
            return_value=model,
        ),
    ):
        return Mask2FormerAdapter._load(
            name="CondadosAI/mask2former_swin_tiny_coco_instance",
            revision="deadbeef",
            device=torch.device("cpu"),
        )


class TestLoad:
    def test_revision_passed_to_both(self) -> None:
        with (
            patch("acaua.adapters.mask2former.AutoImageProcessor.from_pretrained") as proc_from,
            patch(
                "acaua.adapters.mask2former.AutoModelForUniversalSegmentation.from_pretrained"
            ) as model_from,
        ):
            proc_from.return_value = _make_processor_mock()
            model_from.return_value = _make_model_mock()
            Mask2FormerAdapter._load(
                name="CondadosAI/mask2former_swin_tiny_coco_instance",
                revision="v1-sha",
                device=torch.device("cpu"),
            )
            assert proc_from.call_args.kwargs["revision"] == "v1-sha"
            assert model_from.call_args.kwargs["revision"] == "v1-sha"

    def test_eval_and_to_called(self, mocked_adapter: Mask2FormerAdapter) -> None:
        mocked_adapter._model.eval.assert_called_once()
        mocked_adapter._model.to.assert_called_once_with(torch.device("cpu"))


class TestPredict:
    def test_single_input_returns_single_result(
        self, mocked_adapter: Mask2FormerAdapter, rgb_pil
    ) -> None:
        out = mocked_adapter.predict(rgb_pil)
        assert isinstance(out, SegmentationResult)
        assert out.masks.shape[0] == 2  # from the mock

    def test_sequence_input_returns_list(self, mocked_adapter: Mask2FormerAdapter, rgb_pil) -> None:
        out = mocked_adapter.predict([rgb_pil, rgb_pil, rgb_pil])
        assert isinstance(out, list)
        assert len(out) == 3
        assert all(isinstance(r, SegmentationResult) for r in out)

    def test_result_tensor_dtypes(self, mocked_adapter: Mask2FormerAdapter, rgb_pil) -> None:
        out = mocked_adapter.predict(rgb_pil)
        assert isinstance(out, SegmentationResult)
        assert out.boxes.dtype == torch.float32
        assert out.labels.dtype == torch.int64
        assert out.scores.dtype == torch.float32
        assert out.masks.dtype == torch.bool

    def test_batching_chunks_inputs(self, mocked_adapter: Mask2FormerAdapter, rgb_pil) -> None:
        mocked_adapter.predict([rgb_pil] * 7, batch_size=3)
        # 7 inputs, batch_size=3 → 3 batches (3 + 3 + 1).
        assert mocked_adapter._processor.call_count == 3

    def test_forward_under_inference_mode(
        self, mocked_adapter: Mask2FormerAdapter, rgb_pil
    ) -> None:
        mocked_adapter.predict(rgb_pil)
        states = mocked_adapter._model.grad_enabled_during_forward
        assert len(states) >= 1
        assert all(s is False for s in states)


class TestMasksToBoxes:
    def test_basic_box_from_mask(self) -> None:
        mask = torch.zeros(1, 8, 8, dtype=torch.bool)
        mask[0, 2:5, 3:7] = True  # rows 2-4, cols 3-6 → xyxy = [3, 2, 7, 5]
        boxes = _masks_to_boxes(mask)
        assert boxes.shape == (1, 4)
        assert boxes[0].tolist() == [3.0, 2.0, 7.0, 5.0]

    def test_multiple_masks(self) -> None:
        masks = torch.zeros(2, 10, 10, dtype=torch.bool)
        masks[0, 0:3, 0:3] = True  # xyxy [0, 0, 3, 3]
        masks[1, 5:10, 5:10] = True  # xyxy [5, 5, 10, 10]
        boxes = _masks_to_boxes(masks)
        assert boxes.shape == (2, 4)
        assert boxes[0].tolist() == [0.0, 0.0, 3.0, 3.0]
        assert boxes[1].tolist() == [5.0, 5.0, 10.0, 10.0]

    def test_empty_mask_gives_zero_row(self) -> None:
        masks = torch.zeros(1, 8, 8, dtype=torch.bool)
        boxes = _masks_to_boxes(masks)
        assert boxes.shape == (1, 4)
        assert boxes[0].tolist() == [0.0, 0.0, 0.0, 0.0]

    def test_empty_input_returns_zero_rows(self) -> None:
        masks = torch.zeros(0, 8, 8, dtype=torch.bool)
        boxes = _masks_to_boxes(masks)
        assert boxes.shape == (0, 4)

    def test_output_is_float32(self) -> None:
        masks = torch.zeros(1, 8, 8, dtype=torch.bool)
        masks[0, 1:3, 1:3] = True
        boxes = _masks_to_boxes(masks)
        assert boxes.dtype == torch.float32


class TestUnpackInstanceSeg:
    def test_empty_segments_info(self) -> None:
        result = {
            "segmentation": torch.full((16, 16), -1, dtype=torch.int64),
            "segments_info": [],
        }
        out = _unpack_instance_seg(result, torch.device("cpu"))
        assert out.boxes.shape == (0, 4)
        assert out.labels.shape == (0,)
        assert out.scores.shape == (0,)
        assert out.masks.shape == (0, 16, 16)

    def test_extracts_binary_masks_per_instance(self) -> None:
        seg = torch.full((16, 16), -1, dtype=torch.int64)
        seg[0:4, 0:4] = 0  # instance 0
        seg[8:12, 8:12] = 1  # instance 1
        result = {
            "segmentation": seg,
            "segments_info": [
                {"id": 0, "label_id": 5, "score": 0.9},
                {"id": 1, "label_id": 7, "score": 0.8},
            ],
        }
        out = _unpack_instance_seg(result, torch.device("cpu"))
        assert out.masks.shape == (2, 16, 16)
        # Instance 0 mask: pixels in [0:4, 0:4] are True.
        assert out.masks[0, 0:4, 0:4].all()
        assert not out.masks[0, 8:12, 8:12].any()
        # Instance 1 mask: pixels in [8:12, 8:12] are True.
        assert out.masks[1, 8:12, 8:12].all()

    def test_labels_and_scores_come_from_segments_info(self) -> None:
        seg = torch.full((8, 8), -1, dtype=torch.int64)
        seg[0:2, 0:2] = 0
        result = {
            "segmentation": seg,
            "segments_info": [{"id": 0, "label_id": 42, "score": 0.75}],
        }
        out = _unpack_instance_seg(result, torch.device("cpu"))
        assert out.labels.tolist() == [42]
        assert pytest.approx(out.scores.tolist()[0], rel=1e-5) == 0.75

    def test_boxes_match_mask_extents(self) -> None:
        seg = torch.full((16, 16), -1, dtype=torch.int64)
        seg[3:8, 5:12] = 0  # xyxy = [5, 3, 12, 8]
        result = {
            "segmentation": seg,
            "segments_info": [{"id": 0, "label_id": 1, "score": 0.9}],
        }
        out = _unpack_instance_seg(result, torch.device("cpu"))
        assert out.boxes[0].tolist() == [5.0, 3.0, 12.0, 8.0]
