from __future__ import annotations

from pathlib import Path

from adam.image_preferences import GenerationPreferenceEvaluator, PreferenceProfile


def _image(path: Path) -> Path:
    path.write_bytes(b"image")
    return path


def test_generation_ratings_save_load_and_manual_correction(tmp_path: Path) -> None:
    image = _image(tmp_path / "good.png")
    profile = PreferenceProfile(tmp_path, "ddpm_generator", "SML", str(tmp_path / "model"))

    profile.set_rating(image, "keep", seed=12, sampler="DDIM", steps=50)
    loaded = PreferenceProfile(tmp_path, "ddpm_generator", "SML", str(tmp_path / "model"))

    assert loaded.rating_for(image).rating == "keep"  # type: ignore[union-attr]
    loaded.set_rating(image, "reject", seed=12, sampler="DDIM", steps=50)

    corrected = PreferenceProfile(tmp_path, "ddpm_generator", "SML", str(tmp_path / "model"))
    assert corrected.rating_for(image).rating == "reject"  # type: ignore[union-attr]


def test_preference_profiles_are_model_specific(tmp_path: Path) -> None:
    image = _image(tmp_path / "sample.png")
    one = PreferenceProfile(tmp_path, "ddpm_generator", "SML", str(tmp_path / "sml"))
    two = PreferenceProfile(tmp_path, "ddpm_generator", "Minecraft", str(tmp_path / "minecraft"))

    one.set_rating(image, "favorite")

    assert one.id != two.id
    assert PreferenceProfile(tmp_path, "ddpm_generator", "SML", str(tmp_path / "sml")).rating_for(image)
    assert PreferenceProfile(tmp_path, "ddpm_generator", "Minecraft", str(tmp_path / "minecraft")).rating_for(image) is None


def test_preference_scoring_ranks_candidates_with_positive_and_negative_examples(tmp_path: Path) -> None:
    favorite = _image(tmp_path / "favorite.png")
    rejected = _image(tmp_path / "rejected.png")
    close = _image(tmp_path / "close.png")
    far = _image(tmp_path / "far.png")
    vectors = {
        str(favorite): [[1.0, 0.0]],
        str(rejected): [[0.0, 1.0]],
        str(close): [[0.95, 0.05]],
        str(far): [[0.05, 0.95]],
    }
    profile = PreferenceProfile(tmp_path, "flow_generator", "Rooms", str(tmp_path / "flow-model"))
    profile.set_rating(favorite, "favorite")
    profile.set_rating(rejected, "reject")
    evaluator = GenerationPreferenceEvaluator(tmp_path, embedder=lambda paths: [vectors[str(Path(path))][0] for path in paths])

    scores = evaluator.score(profile, [close, far], keep_threshold=0.70, reject_threshold=0.30)

    assert scores[0].score is not None
    assert scores[1].score is not None
    assert scores[0].score > scores[1].score  # type: ignore[operator]
    assert scores[0].category == "Strong Keep"
    assert scores[1].category == "Likely Reject"


def test_scoring_without_profile_signal_needs_review(tmp_path: Path) -> None:
    image = _image(tmp_path / "candidate.png")
    profile = PreferenceProfile(tmp_path, "ddpm_generator", "Empty", str(tmp_path / "empty"))
    evaluator = GenerationPreferenceEvaluator(tmp_path, embedder=lambda _paths: [[1.0, 0.0]])

    scores = evaluator.score(profile, [image])

    assert scores[0].score is None
    assert scores[0].category == "Needs Review"