#!/usr/bin/env python3
"""Test save_state/load_state round trip."""
import zmq, json, time, subprocess, glob, os

ctx = zmq.Context()
cp = ctx.socket(zmq.PUB)
cp.connect("tcp://127.0.0.1:5557")
ack_sub = ctx.socket(zmq.SUB)
ack_sub.connect("tcp://127.0.0.1:5559")
ack_sub.setsockopt_string(zmq.SUBSCRIBE, "")
time.sleep(0.3)

# Test save_state with path "."
print("SAVING to . ...")
cp.send_string(json.dumps({"cmd": "save_state", "path": "."}))
time.sleep(2)
ack_sub.setsockopt(zmq.RCVTIMEO, 5000)
try:
    msg = ack_sub.recv_string()
    print("ACK save:", msg)
except Exception as e:
    print("ERR save:", e)

# Check checkpoint files
files = sorted(glob.glob("/mnt/d/resonance-engine-active/ckpt_*.bin"))
print(f"Ckpt files after save: {len(files)}")
if files:
    latest = files[-1]
    print(f"Latest: {latest}")
    print(f"LOADING {latest} ...")
    cp.send_string(json.dumps({"cmd": "load_state", "path": latest}))
    time.sleep(2)
    ack_sub.setsockopt(zmq.RCVTIMEO, 5000)
    try:
        msg = ack_sub.recv_string()
        print("ACK load:", msg)
    except Exception as e:
        print("ERR load:", e)

cp.close()
ack_sub.close()
ctx.term()
print("DONE")