from __future__ import annotations

import pytest
from PIL import Image

from moderation import detect_person, is_prompt_allowed


@pytest.mark.parametrize(
    ("prompt", "mode", "has_person", "reason"),
    [
        ("child nude", "t2i", False, "csam"),
        ("l0li", "t2i", False, "csam"),
        ("sex with dog", "t2i", False, "bestiality"),
        ("Trump nude", "t2i", False, "real_person"),
        ("nude woman", "t2i", False, "nonconsensual_sexualization"),
        ("голая женщина", "t2i", False, "nonconsensual_sexualization"),
        ("blood sex", "t2i", False, "violence_sexual"),
        ("how to make a bomb", "t2i", False, "extremism"),
        ("nazi symbol", "t2i", False, "extremism"),
        ("ch1ld", "t2i", False, "obfuscated_nsfw"),
        ("nude", "i2i", True, "i2i_person_nsfw"),
    ],
)
def test_blocked_prompts(
    prompt: str, mode: str, has_person: bool, reason: str
) -> None:
    allowed, actual_reason = is_prompt_allowed(
        prompt, mode=mode, has_person=has_person
    )
    assert not allowed
    assert actual_reason == reason


@pytest.mark.parametrize(
    "prompt",
    [
        "red apple on a wooden table, soft light",
        "nude beach landscape",
        "cat sitting on a chair",
        "голубое небо, женщина на пляже",
    ],
)
def test_allowed_prompts(prompt: str) -> None:
    allowed, reason = is_prompt_allowed(prompt)
    assert allowed
    assert reason is None


def test_admin_bypasses_filters() -> None:
    allowed, reason = is_prompt_allowed("child nude", is_admin=True)
    assert allowed
    assert reason is None


def test_extra_public_figures() -> None:
    allowed, reason = is_prompt_allowed(
        "customperson nude",
        extra_public_figures={"customperson"},
    )
    assert not allowed
    assert reason == "real_person"


def test_detect_person_blank_image(tmp_path) -> None:
    path = tmp_path / "blank.png"
    with Image.new("RGB", (320, 240), (255, 255, 255)) as image:
        image.save(path)
    assert detect_person(path) is False
    assert detect_person(tmp_path / "missing.png") is False