"""End-to-end Mask2Former test against the real mirror.

Downloads weights at the pinned SHA, runs instance segmentation on a small
synthetic RGB image, and asserts shape invariants on the returned
`SegmentationResult`. Same no-quality-assertion posture as the RT-DETR E2E.
"""

from __future__ import annotations

import numpy as np
import pytest
import torch
from PIL import Image

import acaua

REPO_ID = "CondadosAI/mask2former_swin_tiny_coco_instance"


def _synthetic_rgb(h: int = 128, w: int = 128) -> Image.Image:
    arr = np.full((h, w, 3), 100, dtype=np.uint8)
    return Image.fromarray(arr, mode="RGB")


@pytest.fixture(scope="module")
def model() -> acaua.Model:
    return acaua.Model.from_pretrained(REPO_ID)


@pytest.mark.e2e
class TestLoad:
    def test_repo_id_stored(self, model: acaua.Model) -> None:
        assert model.repo_id == REPO_ID


@pytest.mark.e2e
class TestPredict:
    def test_single_image_returns_segmentation_result(self, model: acaua.Model) -> None:
        out = model.predict(_synthetic_rgb())
        assert isinstance(out, acaua.SegmentationResult)
        assert out.boxes.ndim == 2 and out.boxes.shape[-1] == 4
        assert out.masks.ndim == 3

    def test_mask_and_box_counts_match(self, model: acaua.Model) -> None:
        """Number of masks must equal number of boxes / labels / scores."""
        out = model.predict(_synthetic_rgb())
        assert isinstance(out, acaua.SegmentationResult)
        n = out.boxes.shape[0]
        assert out.labels.shape[0] == n
        assert out.scores.shape[0] == n
        assert out.masks.shape[0] == n

    def test_batch_returns_list(self, model: acaua.Model) -> None:
        out = model.predict([_synthetic_rgb(), _synthetic_rgb(), _synthetic_rgb()], batch_size=2)
        assert isinstance(out, list)
        assert len(out) == 3

    def test_supervision_conversion_includes_mask(self, model: acaua.Model) -> None:
        import supervision as sv

        out = model.predict(_synthetic_rgb())
        assert isinstance(out, acaua.SegmentationResult)
        det = out.to_supervision()
        assert isinstance(det, sv.Detections)
        # If there are any detections, mask should be populated.
        if len(det) > 0:
            assert det.mask is not None

    def test_mask_dtype_is_bool(self, model: acaua.Model) -> None:
        out = model.predict(_synthetic_rgb())
        assert isinstance(out, acaua.SegmentationResult)
        assert out.masks.dtype == torch.bool
