"""Run the STOCK tinker_cookbook math_rl recipe for one arm, with read-only instrumentation:
- inject Spindle base_url/api_key/Modal proxy headers into tinker.ServiceClient (recipe has no header option);
- before the client's first forward_backward: record get_info (server lora_rank/optimizer_config), fixed-batch
  initial logprobs fingerprint, and measured adapter hash (save_weights_external + download);
- record each optim_step's AdamParams and per-step group reward vectors (from compute_advantages input).
Training logic, loss, optimizer and data order are unchanged."""
import asyncio, json, os, sys, time
import tinker
from tinker import types
import common

tag, rank, lr, steps = sys.argv[1], int(sys.argv[2]), float(sys.argv[3]), int(sys.argv[4])
log_path = f"/home/ubuntu/rx_run_40338c73ff/runs/{tag}"
ev = f"/home/ubuntu/rx_run_40338c73ff/ev/{tag}"
os.makedirs(ev, exist_ok=True)

_orig_init = tinker.ServiceClient.__init__
def _init(self, *a, **kw):
    kw["base_url"] = os.environ["TINKER_BASE_URL"]; kw.setdefault("api_key", os.environ["TINKER_API_KEY"])
    kw["default_headers_overrides"] = common.HEADERS
    _orig_init(self, *a, **kw)
tinker.ServiceClient.__init__ = _init

_orig_create = tinker.ServiceClient.create_lora_training_client_async
async def _create(self, *a, **kw):
    for attempt in range(12):
        try:
            return await _orig_create(self, *a, **kw)
        except Exception as e:
            print(f"[rx] create retry {attempt}: {repr(e)[:300]}", flush=True)
            await asyncio.sleep(min(60 * 2 ** attempt, 240))
    raise RuntimeError("create_lora_training_client failed after retries")
tinker.ServiceClient.create_lora_training_client_async = _create
_orig_restore = tinker.ServiceClient.create_training_client_from_state_with_optimizer_async
async def _restore(self, *a, **kw):
    for attempt in range(12):
        try:
            tc = await _orig_restore(self, *a, **kw)
            json.dump({"restored_from": a[0] if a else kw.get("path"), "time": time.time()}, open(f"{ev}/restore1_request.json", "w"))
            return tc
        except Exception as e:
            print(f"[rx] restore retry {attempt}: {repr(e)[:300]}", flush=True)
            await asyncio.sleep(min(60 * 2 ** attempt, 240))
    raise RuntimeError("restore failed after retries")
tinker.ServiceClient.create_training_client_from_state_with_optimizer_async = _restore

RESUME = os.path.exists(f"{ev}/init_provenance.json")
SUFFIX = "_restore1" if RESUME else ""
state = {"init_done": False, "optim_calls": []}
_orig_fb = tinker.TrainingClient.forward_backward_async
async def _fb(self, *a, **kw):
    if not state["init_done"]:
        state["init_done"] = True
        t0 = time.time(); rec = {"tag": tag, "requested_rank": rank, "lr": lr}
        rec["get_info_before_first_update"] = common.info_dict(await self.get_info_async())
        datums, spec, bsha = common.fixed_batch(self.get_tokenizer())
        fwd = await (await self.forward_async(datums, "cross_entropy"))
        vals, fp = common.fingerprint(fwd)
        rec.update(batch_sha256=bsha, output_sha256=fp, n_logprobs=len(vals))
        json.dump(vals, open(f"{ev}/initial_logprobs{SUFFIX}.json", "w"))
        if os.environ.get("RX_TRY_ADAPTER_HASH") == "1":
            try:
                rest = tinker.ServiceClient().create_rest_client()
                ah = await asyncio.to_thread(common.adapter_hash, self, rest, f"{tag}-init")
                rec["adapter_sha256"] = ah["adapter_sha256"]; rec["adapter_external_path"] = ah["external_path"]
                rec["adapter_files"] = ah["files"]; rec["rank_summary"] = common.rank_summary(ah["tensors"])
                json.dump(ah["tensors"], open(f"{ev}/initial_adapter_tensors.json", "w"))
            except Exception as e:
                rec["adapter_error"] = repr(e)[:500]
        else:
            rec["adapter_error"] = "not attempted: save_weights_external returned 404 Not Found on shared Spindle in the preflight probe"
        rec["instrumentation_seconds"] = time.time() - t0
        rec["is_restore"] = RESUME
        json.dump(rec, open(f"{ev}/init_provenance{SUFFIX}.json", "w"), indent=1, default=str)
        print("[rx] init provenance", json.dumps({k: rec.get(k) for k in ("output_sha256", "batch_sha256", "adapter_error")}), json.dumps(rec["get_info_before_first_update"])[:800], flush=True)
    return await _orig_fb(self, *a, **kw)
tinker.TrainingClient.forward_backward_async = _fb

_orig_os = tinker.TrainingClient.optim_step_async
async def _os(self, adam_params, *a, **kw):
    state["optim_calls"].append(adam_params.model_dump(mode="json") if hasattr(adam_params, "model_dump") else str(adam_params))
    json.dump(state["optim_calls"], open(f"{ev}/optim_calls{SUFFIX}.json", "w"))
    return await _orig_os(self, adam_params, *a, **kw)
tinker.TrainingClient.optim_step_async = _os

import tinker_cookbook.rl.train as rltrain
_orig_adv = rltrain.compute_advantages
def _adv(groups):
    with open(f"{ev}/group_rewards.jsonl", "a") as f:
        f.write(json.dumps([list(map(float, g.get_total_rewards())) for g in groups]) + "\n")
    return _orig_adv(groups)
rltrain.compute_advantages = _adv

import chz
from tinker_cookbook.recipes.math_rl import train as recipe
cfg = chz.entrypoint(recipe.CLIConfig, argv=[
    f"model_name={common.MODEL}", f"renderer_name={common.RENDERER}", f"base_url={os.environ['TINKER_BASE_URL']}",
    "env=gsm8k", f"lora_rank={rank}", "group_size=8", "groups_per_batch=8", f"learning_rate={lr}",
    "max_tokens=512", f"max_steps={steps}", "seed=0", f"log_path={log_path}", "loss_fn=importance_sampling",
    "eval_every=0", "save_every=20", "behavior_if_log_dir_exists=resume"])
json.dump(chz.asdict(cfg) if hasattr(chz, "asdict") else repr(cfg), open(f"{ev}/cli_config.json", "w"), default=str, indent=1)
asyncio.run(recipe.cli_main(cfg))
print("[rx] done", tag, flush=True)