from __future__ import annotations

from io import BytesIO

import numpy as np
import pytest
from PIL import Image

from src.core.services import ml_inference


class _FakeEmbedder:
    loaded = True

    def embed_tiles(self, tiles):
        return np.ones((len(tiles), 384), dtype=np.float32)


class _FakeABMIL:
    loaded = True

    def predict_embeddings(self, embeddings, scale_indices=None):
        assert scale_indices is not None
        assert len(scale_indices) == embeddings.shape[0]
        return {
            "prediction": "Pure",
            "probability_adulterated": 0.2,
            "threshold": 0.5,
            "attention": np.ones((embeddings.shape[0],), dtype=np.float32),
            "num_tiles": int(embeddings.shape[0]),
        }


def _make_png_bytes() -> bytes:
    image = Image.new("RGB", (256, 256), color=(255, 255, 255))
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


@pytest.mark.asyncio
async def test_run_inference_smoke(monkeypatch: pytest.MonkeyPatch) -> None:
    async def _cache_bypass_get(_key: str):
        return None

    monkeypatch.setattr(ml_inference.redis_manager, "get_cache", _cache_bypass_get)

    def _fake_extract_tiles(image_bgr, *, tile_sizes, overlap_ratio, shrink_factor, **kwargs):
        tile_size = tile_sizes[0]
        tile = np.zeros((tile_size, tile_size, 3), dtype=np.uint8)
        return (
            [tile],
            np.asarray([[0, 0]], dtype=np.int32),
            np.asarray([0], dtype=np.int64),
            image_bgr.shape[:2],
        )

    monkeypatch.setattr(ml_inference, "DINO_EMBEDDER", _FakeEmbedder())
    monkeypatch.setattr(ml_inference, "ABMIL_ENGINE", _FakeABMIL())
    monkeypatch.setattr(ml_inference, "extract_tiles_for_inference", _fake_extract_tiles)

    result = await ml_inference.run_inference(_make_png_bytes())
    assert result["predicted_class"] == "Pure Milk"
    assert result["status"] == "success"
    assert isinstance(result["latency_ms"], int)
    assert 0.0 <= result["confidence"] <= 1.0


def test_initialize_models_requires_checkpoint(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(ml_inference, "DINO_EMBEDDER", None)
    monkeypatch.setattr(ml_inference, "ABMIL_ENGINE", None)
    monkeypatch.setattr(ml_inference.settings, "model_checkpoint_path", "")
    monkeypatch.setattr(ml_inference.settings, "abmil_hf_repo_id", None)

    with pytest.raises(RuntimeError, match="ABMIL checkpoint not configured"):
        ml_inference.initialize_models()