from pathlib import Path

import torch
from PIL import Image

from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec
from stimulus_synthesis.scoring import AssetScorer, EncoderPreprocessSpec
from stimulus_synthesis.search.genetic import GeneticSearch
from stimulus_synthesis.spaces import PromptSearchSpace, SeededSearchSpace


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


class SeedColorTextToImage:
    def __init__(self):
        self.calls = []

    def generate(self, prompts, *, generator=None, **kwargs):
        seed = int(generator.initial_seed()) if generator is not None else -1
        self.calls.append((prompts[0], seed))
        value = seed % 256
        return [Image.new("RGB", (8, 8), color=(value, value, value)) for _ in prompts]


class StaticImageToVideo:
    def generate(self, image, prompt, *, generator=None, **kwargs):
        return image


class SeedValueVideo:
    def generate(self, image, prompt, *, generator=None, **kwargs):
        seed = int(generator.initial_seed()) if generator is not None else 0
        value = (seed % 256) / 255.0
        return torch.full((2, 3, 8, 8), value)


def test_seeded_search_space_appends_seed_gene():
    base = PromptSearchSpace(prompt_banks={"subject": ["face"], "style": ["photo"]})
    space = SeededSearchSpace(base, [11, 22])
    cand = space.random_candidate()

    assert len(cand.genes) == 3
    assert space.decode(cand) == "face photo"
    assert space.decode_seed(cand) in {11, 22}
    assert len(space.options) == 3


def test_asset_scored_search_keys_by_prompt_and_seed(tmp_path):
    base = PromptSearchSpace(prompt_banks={"subject": ["same prompt"]})
    space = SeededSearchSpace(base, [0, 255])
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1))
    t2i = SeedColorTextToImage()
    search = GeneticSearch(
        max_evals=2,
        population_size=2,
        n_init=2,
        mutation_rate=1.0,
        image_kwargs={},
        video_kwargs={},
        score_size=8,
        num_frames=1,
        asset_scorer=scorer,
        asset_dir=tmp_path / "candidates",
        asset_type="image",
        image_asset_spec=ImageAssetSpec(width=8, height=8),
    )

    result = search.run(space, t2i, StaticImageToVideo(), MeanScorer(), target=None, seed=123)

    assert result.best_seed == 255
    assert result.best_score == result.best_score_record.score
    assert result.best_export_record.sha256 == result.best_score_record.sha256
    assert result.best_asset_path and Path(result.best_asset_path).exists()
    assert result.best_metadata["score_source"] == "saved_asset"
    assert len({seed for _prompt, seed in t2i.calls}) == 2


def test_video_search_uses_candidate_seed_for_video_generator(tmp_path):
    base = PromptSearchSpace(prompt_banks={"motion": ["move"]})
    space = SeededSearchSpace(base, [0, 127])
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=2))
    fixed_image = Image.new("RGB", (8, 8), color=(0, 0, 0))

    class FixedImageT2I:
        def generate(self, prompts, *, generator=None, **kwargs):
            return [fixed_image for _ in prompts]

    search = GeneticSearch(
        max_evals=2,
        population_size=2,
        n_init=2,
        mutation_rate=1.0,
        image_kwargs={},
        video_kwargs={},
        score_size=8,
        num_frames=2,
        asset_scorer=scorer,
        asset_dir=tmp_path / "videos",
        asset_type="video",
        video_asset_spec=VideoAssetSpec(width=8, height=8, fps=24, num_frames=2, crf=0),
    )

    result = search.run(space, FixedImageT2I(), SeedValueVideo(), MeanScorer(), target=None, seed=123)

    assert result.best_seed == 127
    assert result.best_asset_path and Path(result.best_asset_path).exists()
    assert result.best_export_record.sha256 == result.best_score_record.sha256


def test_same_prompt_seed_and_params_give_same_saved_image_hash(tmp_path):
    base = PromptSearchSpace(prompt_banks={"subject": ["same prompt"]})
    space = SeededSearchSpace(base, [12345])
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1))
    spec = ImageAssetSpec(width=8, height=8)

    def run_once(out_dir):
        search = GeneticSearch(
            max_evals=1,
            population_size=1,
            n_init=1,
            mutation_rate=0.0,
            image_kwargs={},
            video_kwargs={},
            score_size=8,
            num_frames=1,
            asset_scorer=scorer,
            asset_dir=out_dir,
            asset_type="image",
            image_asset_spec=spec,
        )
        return search.run(space, SeedColorTextToImage(), StaticImageToVideo(), MeanScorer(), target=None, seed=999)

    first = run_once(tmp_path / "first")
    second = run_once(tmp_path / "second")

    assert first.best_prompt == second.best_prompt == "same prompt"
    assert first.best_seed == second.best_seed == 12345
    assert first.best_export_record.sha256 == second.best_export_record.sha256
    assert first.best_score_record.sha256 == second.best_score_record.sha256
    assert first.best_score == second.best_score


def test_same_candidate_hash_is_independent_of_evaluation_order(tmp_path):
    base = PromptSearchSpace(prompt_banks={"subject": ["same prompt", "filler prompt"]})
    space = SeededSearchSpace(base, [12345, 54321])
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1))
    spec = ImageAssetSpec(width=8, height=8)

    def run_order(inds, out_dir):
        search = GeneticSearch(
            max_evals=len(inds),
            population_size=len(inds),
            n_init=len(inds),
            mutation_rate=0.0,
            image_kwargs={},
            video_kwargs={},
            score_size=8,
            num_frames=1,
            asset_scorer=scorer,
            asset_dir=out_dir,
            asset_type="image",
            image_asset_spec=spec,
        )
        # Exercise the same evaluation path GeneticSearch uses, but with fixed order.
        from stimulus_synthesis.search.evaluation import EvaluationManager, candidate_key
        evaluator = EvaluationManager()
        search._evaluate(evaluator, space, SeedColorTextToImage(), StaticImageToVideo(), MeanScorer(), None, inds, seed=999)
        key = candidate_key("same prompt", 12345)
        return evaluator.asset_export_cache[key].sha256, evaluator.asset_score_cache[key].score

    first_hash, first_score = run_order([[0, 0], [1, 1]], tmp_path / "first_order")
    second_hash, second_score = run_order([[1, 1], [0, 0]], tmp_path / "second_order")

    assert first_hash == second_hash
    assert first_score == second_score