"""Run one stock math_rl arm on shared Spindle; adds proxy headers and records client-creation evidence."""
import sys, os, json, time, asyncio
from common import *
import tinker, chz

_orig_init = tinker.ServiceClient.__init__
def _init(self, *a, **kw):
    kw["default_headers_overrides"] = HEADERS
    _orig_init(self, *a, **kw)
tinker.ServiceClient.__init__ = _init

ARGS = dict(a.split("=", 1) for a in sys.argv[1:])
LOG = ARGS["log_path"]
os.makedirs(LOG, exist_ok=True)
DATUM, BATCH = fixed_datum()


async def record(kind, path, tc):
    info = (await tc.get_info_async()).model_dump(mode="json")
    from research_exchange.restoration import output_fingerprint
    f = await tc.forward_async([DATUM], "cross_entropy")
    out = await f.result_async()
    lp = [float(x) for x in out.loss_fn_outputs[0]["logprobs"].tolist()]
    rec = {"kind": kind, "requested_path": path, "time": time.time(), "info": info,
           "batch_sha256": sha(BATCH), "output_sha256": output_fingerprint(lp), "logprobs": lp, "before_first_update": True}
    with open(os.path.join(LOG, "client_events.jsonl"), "a") as fh:
        fh.write(json.dumps(rec) + "\n")
    print("CLIENT_EVENT", kind, info.get("lora_rank"), rec["output_sha256"][:12], flush=True)


for meth, kind in (("create_training_client_from_state_async", "init_load"),
                   ("create_training_client_from_state_with_optimizer_async", "resume_restore"),
                   ("create_lora_training_client_async", "fresh_create")):
    orig = getattr(tinker.ServiceClient, meth)
    def make(orig, kind):
        async def wrapped(self, *a, **kw):
            tc = await orig(self, *a, **kw)
            await record(kind, a[0] if a else None, tc)
            return tc
        return wrapped
    setattr(tinker.ServiceClient, meth, make(orig, kind))

from tinker_cookbook.recipes.math_rl.train import CLIConfig, cli_main
cfg = chz.entrypoint(CLIConfig, argv=sys.argv[1:])
for attempt in range(8):
    try:
        asyncio.run(cli_main(cfg))
        break
    except Exception as e:
        print("ARM_FAIL attempt", attempt, repr(e)[:400], "cause", repr(e.__cause__)[:300], flush=True)
        time.sleep(min(240, 30 * 2 ** attempt))
print("ARM_DONE", flush=True)