"""Spec preflight probe + restoration preflight fixtures (run before any arm trains)."""
import json, time, urllib.request, hashlib, traceback
import tinker, logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
for n in ("httpx","httpcore","urllib3","filelock","huggingface_hub"): logging.getLogger(n).setLevel(logging.WARNING)
from common import *

t0 = time.time()
datum, batch = fixed_datum()
batch_sha = sha(batch)
json.dump(batch, open("fixed_batch.json", "w"))
sc = service_client()
rest = sc.create_rest_client()
res = {"model": MODEL, "renderer": RENDERER, "batch_sha256": batch_sha, "ranks": {}}


def create(rank):
    for attempt in range(12):
        try:
            return sc.create_lora_training_client(MODEL, rank=rank)
        except Exception as e:
            print("create failed", rank, repr(e)[:300], flush=True)
            time.sleep(min(60, 10 * 2 ** attempt))
    raise RuntimeError("capacity")


for R in (1, 32):
    r = {}
    ts = time.time()
    tc = create(R)
    r["create_seconds"] = time.time() - ts
    r["info_after_create"] = info_dump(tc)
    lp, fp = forward_fingerprint(tc, datum)
    r["initial_output_sha256"] = fp
    r["initial_logprobs_head"] = lp[:5]
    r["initial_wnll"] = -sum(a * w for a, w in zip(lp, batch["weights"])) / sum(batch["weights"])
    st = tc.save_state(f"{RUN}-init-r{R}").result()
    r["init_state_path"] = st.path
    r["adapter_sha256"] = None  # save_weights_external raised and then wedged the client in attempt 2; identity proven by immutable init checkpoint + output fingerprints
    # one update on the probe client
    try:
        fb = tc.forward_backward([datum], "cross_entropy").result(timeout=900)
        r["fb_loss_metrics"] = fb.metrics
        os_ = tc.optim_step(tinker.AdamParams(learning_rate=1e-3, beta1=0.9, beta2=0.95, eps=1e-8)).result()
        r["optim_step_response"] = os_.model_dump(mode="json")
        r["info_after_update"] = info_dump(tc)
        lp3, fp3 = forward_fingerprint(tc, datum)
        r["post_update_wnll"] = -sum(a * w for a, w in zip(lp3, batch["weights"])) / sum(batch["weights"])
        r["update_ok"] = True
    except Exception as e:
        r["update_ok"] = False
        r["update_error"] = repr(e)[:800] + " cause=" + repr(getattr(e, "__cause__", None))[:500]
    # restore into a new client exactly as the arms will, and re-fingerprint before any update
    ts = time.time()
    sc.close('success').result()
    sc2 = service_client()
    tc2 = sc2.create_training_client_from_state(st.path)
    r["restore_seconds"] = time.time() - ts
    r["info_after_restore"] = info_dump(tc2)
    lp2, fp2 = forward_fingerprint(tc2, datum)
    r["restored_output_sha256"] = fp2
    r["restored_max_abs_diff"] = max(abs(a - b) for a, b in zip(lp, lp2))
    sc2.close('success').result()
    sc = service_client()
    res["ranks"][str(R)] = r
    print(json.dumps({k: v for k, v in r.items() if k not in ("external_files",)}, default=str)[:3000], flush=True)

res["fresh_outputs_identical_across_ranks"] = res["ranks"]["1"]["initial_output_sha256"] == res["ranks"]["32"]["initial_output_sha256"]
res["seconds"] = time.time() - t0
json.dump(res, open("probe_results.json", "w"), indent=1, default=str)
print("DONE", res["fresh_outputs_identical_across_ranks"], res["seconds"])