import json

import torch
from PIL import Image

from stimulus_synthesis.asset_manifest import load_asset_manifest, write_asset_manifest
from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec, decode_image, decode_video, export_image, export_video
from stimulus_synthesis.scoring import AssetScorer, EncoderPreprocessSpec, prepare_video_for_encoder


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


def test_image_export_decode_and_asset_score(tmp_path):
    image = Image.new("RGB", (6, 6), color=(128, 64, 32))
    spec = ImageAssetSpec(width=12, height=10, format="png")
    path = tmp_path / "stimulus.png"

    export_record = export_image(image, path, spec)
    decoded = decode_image(path)
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=3))
    score_record = scorer.score_image(path, asset_spec=spec, metadata={"prompt": "test prompt"})

    assert export_record.sha256 == decoded.sha256 == score_record.sha256
    assert decoded.image.shape == (3, 10, 12)
    assert score_record.sampled_frame_indices == [0, 0, 0]
    assert score_record.preprocess["size"] == 8
    assert score_record.asset_spec["width"] == 12
    assert isinstance(score_record.score, float)


def test_video_export_decode_preprocess_and_manifest(tmp_path):
    frames = torch.zeros(3, 3, 8, 8)
    frames[1] = 0.5
    frames[2] = 1.0
    spec = VideoAssetSpec(width=16, height=16, fps=24, num_frames=5, crf=18)
    path = tmp_path / "stimulus.mp4"

    export_record = export_video(frames, path, spec)
    decoded = decode_video(path)
    prepared = prepare_video_for_encoder(decoded.frames, EncoderPreprocessSpec(size=(8, 8), num_frames=4))
    scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=(8, 8), num_frames=4))
    score_record = scorer.score_video(path, asset_spec=spec)

    assert export_record.sha256 == decoded.sha256 == score_record.sha256
    assert decoded.frames.shape[1:] == (3, 16, 16)
    assert decoded.num_frames == 5
    assert prepared.videos.shape == (1, 4, 3, 8, 8)
    assert score_record.sampled_frame_indices == [0, 1, 3, 4]

    manifest_path = tmp_path / "manifest.json"
    manifest = write_asset_manifest([export_record, score_record], manifest_path, metadata={"model": "mock"})
    loaded = load_asset_manifest(manifest_path)
    assert loaded == manifest
    assert loaded["metadata"]["model"] == "mock"
    assert len(loaded["records"]) == 2