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 (
    _compute_frame_rows_independent,
    _compute_frame_rows_sequential,
    _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)
    compute_rows = (
        _compute_frame_rows_independent
        if args.independent_replay
        else _compute_frame_rows_sequential
    )
    rows = compute_rows(
        episode_dir=episode_dir,
        demo=demo,
        templates=templates,
        checkpoint_stride=args.checkpoint_stride,
        frame_indices=args.frame_indices,
    )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    for row in rows:
        frame_index = int(row["frame_index"])
        with output_dir.joinpath(f"frame_{frame_index:04d}.json").open(
            "w", encoding="utf-8"
        ) as handle:
            json.dump(row, handle)
    return 0


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