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

Network-free. Verifies:
  - `from_pretrained` passes the pinned `revision` to the Auto classes.
  - `.eval()` and `.to(device)` are applied after load.
  - Forward pass runs inside `torch.inference_mode()` (grad disabled).
  - Single input returns a single `DetectionResult`; sequence input returns
    `list[DetectionResult]`.
  - Batching chunks a long input list into `batch_size`-sized groups.
  - Output tensors are correctly typed (float32 for boxes/scores, int64 for labels).
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
import torch

from acaua import DetectionResult
from acaua.adapters.rtdetr import RTDETRAdapter


def _make_processor_mock(n_detections_per_image: int = 2) -> MagicMock:
    """Mock AutoImageProcessor instance."""
    processor = MagicMock()

    def call(images, return_tensors):  # type: ignore[no-untyped-def]
        # Pretend to produce a batch dict with pixel_values [B, 3, H, W].
        b = len(images)
        return {"pixel_values": torch.zeros(b, 3, 32, 32)}

    processor.side_effect = call

    def post_process(outputs, target_sizes, threshold):  # type: ignore[no-untyped-def]
        # Return one dict per image with plausible tensors.
        return [
            {
                "boxes": torch.zeros(n_detections_per_image, 4),
                "labels": torch.zeros(n_detections_per_image, dtype=torch.int64),
                "scores": torch.full((n_detections_per_image,), 0.9),
            }
            for _ in target_sizes
        ]

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


def _make_model_mock() -> MagicMock:
    """Mock the HF model. Records `torch.is_grad_enabled()` during forward."""
    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() -> RTDETRAdapter:
    processor = _make_processor_mock()
    model = _make_model_mock()
    with (
        patch(
            "acaua.adapters.rtdetr.AutoImageProcessor.from_pretrained",
            return_value=processor,
        ) as proc_from,
        patch(
            "acaua.adapters.rtdetr.AutoModelForObjectDetection.from_pretrained",
            return_value=model,
        ) as model_from,
    ):
        adapter = RTDETRAdapter._load(
            name="CondadosAI/rtdetr_r50vd",
            revision="cafef00d",
            device=torch.device("cpu"),
        )
        # Store the patches' call args for inspection in the load-path tests.
        adapter._test_proc_from = proc_from  # type: ignore[attr-defined]
        adapter._test_model_from = model_from  # type: ignore[attr-defined]
    return adapter


class TestLoad:
    def test_passes_revision_to_processor(self, mocked_adapter: RTDETRAdapter) -> None:
        mocked_adapter._test_proc_from.assert_called_once()  # type: ignore[attr-defined]
        _, kwargs = mocked_adapter._test_proc_from.call_args  # type: ignore[attr-defined]
        assert kwargs["revision"] == "cafef00d"

    def test_passes_revision_to_model(self, mocked_adapter: RTDETRAdapter) -> None:
        mocked_adapter._test_model_from.assert_called_once()  # type: ignore[attr-defined]
        _, kwargs = mocked_adapter._test_model_from.call_args  # type: ignore[attr-defined]
        assert kwargs["revision"] == "cafef00d"

    def test_passes_repo_id_to_both(self, mocked_adapter: RTDETRAdapter) -> None:
        p_args, _ = mocked_adapter._test_proc_from.call_args  # type: ignore[attr-defined]
        m_args, _ = mocked_adapter._test_model_from.call_args  # type: ignore[attr-defined]
        assert p_args[0] == "CondadosAI/rtdetr_r50vd"
        assert m_args[0] == "CondadosAI/rtdetr_r50vd"

    def test_eval_called(self, mocked_adapter: RTDETRAdapter) -> None:
        mocked_adapter._model.eval.assert_called_once()

    def test_to_device_called(self, mocked_adapter: RTDETRAdapter) -> None:
        mocked_adapter._model.to.assert_called_once_with(torch.device("cpu"))

    def test_repo_id_stored(self, mocked_adapter: RTDETRAdapter) -> None:
        assert mocked_adapter.repo_id == "CondadosAI/rtdetr_r50vd"

    def test_device_stored(self, mocked_adapter: RTDETRAdapter) -> None:
        assert mocked_adapter.device == torch.device("cpu")


class TestPredictReturnShape:
    def test_single_input_returns_single_result(
        self, mocked_adapter: RTDETRAdapter, rgb_pil
    ) -> None:
        out = mocked_adapter.predict(rgb_pil)
        assert isinstance(out, DetectionResult)

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

    def test_output_tensor_dtypes(self, mocked_adapter: RTDETRAdapter, rgb_pil) -> None:
        out = mocked_adapter.predict(rgb_pil)
        assert isinstance(out, DetectionResult)
        assert out.boxes.dtype == torch.float32
        assert out.labels.dtype == torch.int64
        assert out.scores.dtype == torch.float32


class TestBatching:
    def test_batch_size_1(self, mocked_adapter: RTDETRAdapter, rgb_pil) -> None:
        inputs = [rgb_pil] * 5
        out = mocked_adapter.predict(inputs, batch_size=1)
        assert isinstance(out, list)
        assert len(out) == 5
        # Processor should have been called 5 times (once per batch).
        assert mocked_adapter._processor.call_count == 5

    def test_batch_size_2_on_5_inputs(self, mocked_adapter: RTDETRAdapter, rgb_pil) -> None:
        inputs = [rgb_pil] * 5
        out = mocked_adapter.predict(inputs, batch_size=2)
        assert isinstance(out, list)
        assert len(out) == 5
        # 5 inputs, batch_size=2 → batches of 2, 2, 1 → 3 processor calls.
        assert mocked_adapter._processor.call_count == 3

    def test_default_batch_size_8(self, mocked_adapter: RTDETRAdapter, rgb_pil) -> None:
        out = mocked_adapter.predict([rgb_pil] * 10)
        assert isinstance(out, list)
        assert len(out) == 10
        # 10 inputs, default batch_size=8 → 2 batches (8 + 2).
        assert mocked_adapter._processor.call_count == 2


class TestInferenceMode:
    def test_forward_runs_with_grad_disabled(self, mocked_adapter: RTDETRAdapter, rgb_pil) -> None:
        mocked_adapter.predict(rgb_pil)
        grad_states = mocked_adapter._model.grad_enabled_during_forward
        assert len(grad_states) >= 1
        assert all(state is False for state in grad_states), (
            f"forward pass ran with grad enabled: {grad_states}"
        )
