"""Tests for inference validation and device selection."""

from pathlib import Path
from types import SimpleNamespace
import tempfile
import unittest
from unittest.mock import patch

from PIL import Image
import torch

from rmbg_mask.inference import (
    RmbgError,
    inferMask,
    loadModel,
    selectDevice,
    validateModelPath,
)


class Availability:
    """Provide a configurable is_available method for device tests."""

    def __init__(self, available: bool) -> None:
        """Store the availability result."""
        self.available = available

    def is_available(self) -> bool:
        """Return the configured availability result."""
        return self.available


class FakeModel:
    """Return predictable logits while recording the model input shape."""

    def __init__(self) -> None:
        """Initialize without a recorded input shape."""
        self.input_shape: tuple[int, ...] | None = None

    def __call__(self, input_tensor: torch.Tensor) -> list[torch.Tensor]:
        """Return a gradient of logits matching the input dimensions."""
        self.input_shape = tuple(input_tensor.shape)
        height, width = input_tensor.shape[-2:]
        logits = torch.linspace(-4, 4, height * width)
        return [logits.reshape(1, 1, height, width)]


class LoadedModel:
    """Record evaluation and device placement by the model loader."""

    def __init__(self) -> None:
        """Initialize model state flags."""
        self.evaluated = False
        self.device: str | None = None

    def eval(self) -> "LoadedModel":
        """Record evaluation mode and return this model."""
        self.evaluated = True
        return self

    def to(self, device: str) -> "LoadedModel":
        """Record device placement and return this model."""
        self.device = device
        return self


class InferenceTest(unittest.TestCase):
    """Verify local model validation and automatic device selection."""

    def testModelDirectoryRequiresConfig(self) -> None:
        """Reject a directory that is not a complete model directory."""
        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaisesRegex(RmbgError, "config.json"):
                validateModelPath(Path(directory))

    def testModelDirectoryIsResolved(self) -> None:
        """Resolve a model directory containing its configuration."""
        with tempfile.TemporaryDirectory() as directory:
            model_path = Path(directory)
            (model_path / "config.json").touch()
            validated_path = validateModelPath(model_path)
            self.assertEqual(validated_path, model_path.resolve())

    def testAutoDevicePrefersCuda(self) -> None:
        """Prefer CUDA over other available accelerators."""
        torch_module = SimpleNamespace(
            cuda=Availability(True),
            backends=SimpleNamespace(mps=Availability(True)),
        )
        self.assertEqual(selectDevice(torch_module, "auto"), "cuda")

    def testAutoDeviceFallsBackToCpu(self) -> None:
        """Use CPU when no supported accelerator is available."""
        torch_module = SimpleNamespace(
            cuda=Availability(False),
            backends=SimpleNamespace(mps=Availability(False)),
        )
        self.assertEqual(selectDevice(torch_module, "auto"), "cpu")

    def testInferenceRestoresArbitraryImageSize(self) -> None:
        """Return a soft grayscale mask at the original image dimensions."""
        image = Image.new("RGB", (37, 19), "white")
        model = FakeModel()

        mask = inferMask(model, image, 64, "cpu")

        self.assertEqual(model.input_shape, (1, 3, 64, 64))
        self.assertEqual(mask.mode, "L")
        self.assertEqual(mask.size, image.size)
        populated_values = sum(count > 0 for count in mask.histogram())
        self.assertGreater(populated_values, 2)

    def testModelLoadingIsStrictlyLocal(self) -> None:
        """Tell Transformers to load custom code and weights locally only."""
        loaded_model = LoadedModel()
        auto_model = SimpleNamespace()
        auto_model.from_pretrained = unittest.mock.Mock(
            return_value=loaded_model
        )
        transformers = SimpleNamespace(
            AutoModelForImageSegmentation=auto_model
        )

        with patch.dict("sys.modules", {"transformers": transformers}):
            result = loadModel(Path("/local/model"), "cpu")

        auto_model.from_pretrained.assert_called_once_with(
            "/local/model",
            local_files_only=True,
            trust_remote_code=True,
        )
        self.assertIs(result, loaded_model)
        self.assertTrue(loaded_model.evaluated)
        self.assertEqual(loaded_model.device, "cpu")


if __name__ == "__main__":
    unittest.main()