from pathlib import Path
import argparse
import json
import pickle
import sys


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from rr_label_study.oven_study import (
    BimanualTakeTrayOutOfOven,
    ReplayCache,
    _build_frame_artifacts,
    _launch_replay_env,
    _load_demo,
)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--episode-dir", required=True)
    parser.add_argument("--templates-pkl", required=True)
    parser.add_argument("--frame-indices", nargs="+", type=int, required=True)
    parser.add_argument("--checkpoint-stride", type=int, default=16)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--independent-replay", action="store_true")
    args = parser.parse_args()

    episode_dir = Path(args.episode_dir)
    with Path(args.templates_pkl).open("rb") as handle:
        templates = pickle.load(handle)
    demo = _load_demo(episode_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    env = _launch_replay_env()
    try:
        task = env.get_task(BimanualTakeTrayOutOfOven)
        cache = ReplayCache(task, demo, checkpoint_stride=args.checkpoint_stride)
        cache.reset()
        initial_snapshot = cache.snapshot() if args.independent_replay else None
        for frame_index in sorted({int(index) for index in args.frame_indices}):
            if args.independent_replay:
                cache.restore_to_index(initial_snapshot, 0)
            cache.step_to(frame_index)
            frame_snapshot = cache.snapshot() if not args.independent_replay else None
            state = cache.current_state()
            row, debug = _build_frame_artifacts(
                episode_dir=episode_dir,
                demo=demo,
                templates=templates,
                task=task,
                state=state,
            )
            with output_dir.joinpath(f"frame_{frame_index:04d}.json").open(
                "w", encoding="utf-8"
            ) as handle:
                json.dump(row, handle)
            with output_dir.joinpath(f"frame_{frame_index:04d}.debug.json").open(
                "w", encoding="utf-8"
            ) as handle:
                json.dump(debug, handle)
            if frame_snapshot is not None:
                cache.restore(frame_snapshot)
    finally:
        env.shutdown()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())