"""Replay Nodecode's ARC-AGI-3 games and check them against games.csv.

ARC-AGI-3 games are deterministic: the same actions from a fresh game give the same frames, levels and score. This
script makes each of the 25 public games fresh in ARC Prize's own toolkit (arc-agi, offline, seed 0), plays the
actions in actions/<game>.jsonl one by one, and compares the end state, levels, action count, the toolkit's
scorecard score and a SHA-256 of the final frame with games.csv.

  pip install arc-agi==0.9.9
  python verify.py --games path/to/environment_files

--games is the toolkit's environment_files directory holding the 25 public games. In its default mode the toolkit
downloads a game there the first time Arcade().make("<game_id>") runs (an ARC_API_KEY from arcprize.org); after that
this script replays offline. The game ids, with their versions, are the game_id column of games.csv.
"""
import argparse, csv, hashlib, json, os, sys, tempfile

from arc_agi import Arcade, OperationMode
from arcengine import GameAction

ap = argparse.ArgumentParser()
ap.add_argument("--games", required=True, help="the toolkit's environment_files directory (one folder per game)")
ap.add_argument("--only", nargs="*", help="replay only these games, e.g. ar25 sk48")
args = ap.parse_args()
here = os.path.dirname(os.path.abspath(__file__))
want = {r["game"]: r for r in csv.DictReader(open(os.path.join(here, "games.csv")))}
os.chdir(tempfile.mkdtemp())  # the toolkit writes its own files here, never next to the log

bad = 0
for g in sorted(args.only or want):
    w = want[g]
    arcade = Arcade(operation_mode=OperationMode.OFFLINE, environments_dir=args.games)
    env = arcade.make(w["game_id"], seed=0, save_recording=False, include_frame_data=True)
    last = env.reset()
    n = 0
    for line in open(os.path.join(here, "actions", g + ".jsonl")):
        a = json.loads(line)
        if a["action"] == "RESET":
            last = env.reset()
        elif a["action"] == "ACTION6":
            last = env.step(GameAction.ACTION6, data={"x": a["x"], "y": a["y"]})
        else:
            last = env.step(getattr(GameAction, a["action"]))
        n += 1
        if last.levels_completed != a["levels"]:
            print(f"{g}: action {n} reached {last.levels_completed} levels, the log says {a['levels']}")
            bad += 1
            break
    frame = [z.tolist() if hasattr(z, "tolist") else z for z in last.frame]
    card = json.loads(arcade.get_scorecard().model_dump_json())
    got = dict(state=last.state.name, levels=str(last.levels_completed), actions=str(n), score=f"{float(card['score']):.2f}",
               final_frame_sha256=hashlib.sha256(json.dumps(frame).encode()).hexdigest())
    exp = dict(state=w["state"], levels=w["levels"], actions=w["actions"], score=f"{float(w['score']):.2f}", final_frame_sha256=w["final_frame_sha256"])
    diff = {k: (got[k], exp[k]) for k in got if got[k] != exp[k]}
    bad += bool(diff)
    print(f"{g}  {got['state']:<12} levels {got['levels']:>2}/{w['total_levels']:<2} actions {got['actions']:>4}  score {got['score']:>6}  "
          + ("ok" if not diff else f"MISMATCH {diff}"))
print("all games replay to the logged result" if not bad else f"{bad} game(s) differ")
sys.exit(1 if bad else 0)