import torch
from PIL import Image

from stimulus_synthesis import NevoPipeline
from stimulus_synthesis.spaces import StructuredArtPromptSpace, PromptSearchSpace


class MockTextToImage:
    def generate(self, prompts, *, generator=None, **kwargs):
        return [Image.new("RGB", (8, 8)) for _ in prompts]


class MockImageToVideo:
    def generate(self, image, prompt, *, generator=None, **kwargs):
        return torch.full((2, 3, 8, 8), len(prompt) / 100.0)

    def generate_batch(self, images, prompts, *, generators=None, **kwargs):
        return [self.generate(i, p) for i, p in zip(images, prompts)]


class MockScorer:
    def score(self, videos, target, **kwargs):
        return videos.mean(dim=(1, 2, 3, 4)).tolist()


def _pipe():
    return NevoPipeline(text_to_image=MockTextToImage(), image_to_video=MockImageToVideo(), scorer=MockScorer())


def test_roi_selects_enhanced_structured_space():
    pipe = _pipe()
    enhanced = pipe.make_search_space(roi="FFA")
    general = pipe.make_search_space(roi="FFA", enforce_general_search_space=True)
    assert isinstance(enhanced, StructuredArtPromptSpace)
    assert isinstance(general, StructuredArtPromptSpace)
    # ROI-enhanced space searches strictly fewer categories than the general one
    assert sum(enhanced.active_mask) < sum(general.active_mask)


def test_no_roi_uses_general_prompt_space():
    assert isinstance(_pipe().make_search_space(seed_prompts=["a cat"]), PromptSearchSpace)


def test_pipeline_runs_with_roi_and_general_override():
    pipe = _pipe()
    out = pipe(roi="FFA", image_max_evals=2, video_max_evals=2, population_size=2, seed=0, score_size=None)
    assert out.best_prompt and isinstance(out.best_score, float)
    out2 = pipe(roi="FFA", enforce_general_search_space=True, image_max_evals=2, video_max_evals=2, population_size=2, seed=0, score_size=None)
    assert out2.best_prompt is not None