from pathlib import Path

import torch
from PIL import Image
from diffusers import DiffusionPipeline

from stimulus_synthesis.pipeline import NevoPipeline
from stimulus_synthesis.spaces import PromptSearchSpace


class MockTextToImage:
    def generate(self, prompts, *, generator=None, **kwargs):
        return [Image.new("RGB", (8, 8), color=(len(prompt) % 255, 0, 0)) for prompt in prompts]


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

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


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


def test_pipeline_smoke_with_mock_components():
    space = PromptSearchSpace(
        prompt_banks={
            "subject": ["a", "long subject"],
            "action": ["runs", "walks"],
        }
    )
    pipe = NevoPipeline(
        text_to_image=MockTextToImage(),
        image_to_video=MockImageToVideo(),
        scorer=MockScorer(),
        search_space=space,
    )

    out = pipe(target=[0], image_max_evals=2, population_size=2, seed=123, score_size=None)

    assert out.best_prompt
    assert isinstance(out.best_score, float)
    assert out.best.image is not None
    assert out.best.video is not None
    assert out.metadata["max_evals"] == 2


def test_diffusers_custom_pipeline_loads_local_repo():
    repo = str(Path(__file__).resolve().parents[1])
    pipe = DiffusionPipeline.from_pretrained(
        repo,
        custom_pipeline=repo,
    )

    assert isinstance(pipe, NevoPipeline)
    assert pipe.synthesis_config.encoder_call == "predict_fmri"