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

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


def test_config():
    assert MODEL_CONFIG["architecture"] == "convnext_base"
    assert MODEL_CONFIG["num_params"] == 89_000_000
    assert MODEL_CONFIG["input_size"] == 224


def test_architecture_output_shape():
    model = ImagerPro(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 = ImagerPro(pretrained=False)
    n = sum(p.numel() for p in model.parameters())
    # ConvNeXt-Base has ~88.6M params
    assert 85_000_000 < n < 92_000_000, f"Unexpected param count: {n}"


@pytest.mark.skipif(
    not os.path.exists("model/model_imager-1-pro.pth"),
    reason="Weights not present",
)
def test_predict_no_nan():
    """Critical test: model must not output NaN."""
    model = load_model("model/model_imager-1-pro.pth")
    img = Image.new("RGB", (512, 512), color=(128, 128, 128))
    result = predict(img, model=model)
    assert result["probability"] == result["probability"], "Probability is NaN!"
    assert 0.0 <= result["probability"] <= 1.0
    assert result["label"] in ("Real", "AI-Generated")