"""Export recorded RoboLab comparisons; no new rollouts or generated imagery.

Usage: python3 export.py /path/to/vast-robolab-b-20261024/results/ft-spatula-01/eval
Requires ffmpeg and ffprobe. Output is written beside this script.
"""
import argparse
import hashlib
import json
from pathlib import Path
import subprocess


def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def probe(path):
    return json.loads(subprocess.check_output([
        "ffprobe", "-v", "error", "-show_entries",
        "stream=width,height,r_frame_rate:format=duration", "-of", "json", str(path)
    ], text=True))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("evaluation_root", type=Path)
    args = parser.parse_args()
    output = Path(__file__).resolve().parent
    media = output / "media"
    media.mkdir(exist_ok=True)
    selection = {"run": 1, "env_id": 8}
    provenance = {
        "study": "ft-spatula-01",
        "task": "WoodSpatulaToBowlTask",
        "instruction": "Put the wooden spatula in the bowl",
        "selection": selection,
        "selection_note": "Selected after inspecting outcomes and recordings: base succeeds in PhysX and fails in Newton; fine-tuned succeeds in both. Illustrative, not an average outcome.",
        "pairing": "Same run and environment index; paired policy sampling noise. Newton execution is not repeatable across reruns.",
        "rendering": "Original Isaac Sim RTX egocentric camera recordings at 15 fps. PhysX left, Newton right. Both start at the first recorded frame and play at original speed. The shorter episode holds its final frame until the longer ends. No trajectories are reconstructed.",
        "sources": [],
        "outputs": [],
    }
    episodes = {}
    for policy in ["base", "finetuned"]:
        inputs, durations = [], []
        for engine in ["physx", "newton_mjwarp"]:
            cell = f"ft-spatula-01-eval-{policy}-{engine}"
            root = args.evaluation_root / cell
            rows = [json.loads(line) for line in (root / "episode_results.jsonl").read_text().splitlines()]
            fields = ["run", "env_id", "success", "duration", "events", "physics_artifact", "reason"]
            episodes[cell] = [{key: row[key] for key in fields} for row in rows]
            selected = next(row for row in rows if all(row[key] == value for key, value in selection.items()))
            video = root / "WoodSpatulaToBowlTask/Put_the_wooden_spatula_in_the_bowl_1_env8_viewport.mp4"
            metadata = probe(video)
            duration = float(metadata["format"]["duration"])
            inputs.append(video)
            durations.append(duration)
            provenance["sources"].append({
                "policy": policy, "engine": engine,
                "path": str(video.relative_to(args.evaluation_root)),
                "sha256": sha256(video), "video": metadata,
                "episode": {key: selected[key] for key in fields},
            })
        duration = max(durations)
        # Keep the original pixels and timebase; only compose and compress for the page.
        filters = []
        for index, length in enumerate(durations):
            filters.append(f"[{index}:v]setpts=PTS-STARTPTS,tpad=stop_mode=clone:stop_duration={duration-length:.6f}[v{index}]")
        filters.append("[v0][v1]hstack=inputs=2:shortest=1[v]")
        video = media / f"spatula-{policy}-paired.mp4"
        subprocess.run([
            "ffmpeg", "-v", "error", "-y", "-i", str(inputs[0]), "-i", str(inputs[1]),
            "-filter_complex", ";".join(filters), "-map", "[v]", "-an",
            "-c:v", "libx264", "-preset", "medium", "-crf", "23", "-pix_fmt", "yuv420p",
            "-movflags", "+faststart", str(video),
        ], check=True)
        poster = video.with_suffix(".jpg")
        subprocess.run([
            "ffmpeg", "-v", "error", "-y", "-ss", "1", "-i", str(video),
            "-frames:v", "1", "-q:v", "3", str(poster),
        ], check=True)
        for path in [video, poster]:
            provenance["outputs"].append({"path": str(path.relative_to(output)), "sha256": sha256(path), "bytes": path.stat().st_size})
    (output / "episode-results.json").write_text(json.dumps(episodes, indent=2) + "\n")
    (output / "media-provenance.json").write_text(json.dumps(provenance, indent=2) + "\n")
    print(json.dumps({"outputs": provenance["outputs"], "episodes": {k: len(v) for k, v in episodes.items()}}, indent=2))


if __name__ == "__main__":
    main()