"""Unit tests for Imager-1-Tiny."""
import os
import torch
import pytest
from PIL import Image

from src.config import ImagerTiny, MODEL_CONFIG, load_model, predict


def test_config_threshold():
    assert MODEL_CONFIG["threshold"] == 0.385
    assert MODEL_CONFIG["input_size"] == 224
    assert MODEL_CONFIG["num_classes"] == 1


def test_architecture_output_shape():
    model = ImagerTiny(pretrained=False)
    model.eval()
    dummy = torch.randn(1, 3, 224, 224)
    with torch.no_grad():
        out = model(dummy)
    assert out.shape == (1, 1), f"Expected (1,1), got {out.shape}"


def test_architecture_param_count():
    model = ImagerTiny(pretrained=False)
    n = sum(p.numel() for p in model.parameters())
    # ResNet-50 has ~25.6M params; allow a margin
    assert 24_000_000 < n < 27_000_000, f"Unexpected param count: {n}"


@pytest.mark.skipif(
    not os.path.exists("model/model_imager-1-tiny.pth"),
    reason="Weights not downloaded",
)
def test_predict_runs():
    model = load_model("model/model_imager-1-tiny.pth")
    img = Image.new("RGB", (512, 512), color=(128, 128, 128))
    result = predict(img, model=model)
    assert "probability" in result
    assert result["label"] in ("Real", "AI-Generated")
    assert 0.0 <= result["probability"] <= 1.0