"""Shared helpers for run_40338c73ff: Spindle client, fixed batch, adapter hashing, frozen GSM8K evaluator."""
import hashlib, io, json, math, os, re, urllib.request
import tinker
from tinker import types

MODEL = "Qwen/Qwen2.5-3B"
RENDERER = "role_colon"  # cookbook model_info has no Qwen2.5-3B entry; role_colon is its base-model renderer (Llama-3.1-8B base)
HEADERS = {"Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"], "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"]}

def service():
    return tinker.ServiceClient(base_url=os.environ["TINKER_BASE_URL"], api_key=os.environ["TINKER_API_KEY"],
                                default_headers_overrides=HEADERS)

def sha(b):
    return hashlib.sha256(b).hexdigest()

def jdump(obj):
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()

def info_dict(info):
    d = info.model_dump(mode="json") if hasattr(info, "model_dump") else dict(info)
    return d

def fixed_batch(tokenizer):
    """Fixed tokenized batch: GSM8K train[0] question + reference solution through role_colon renderer."""
    from datasets import load_dataset
    from tinker_cookbook import renderers
    from tinker_cookbook.recipes.math_rl.math_env import MathEnv
    r = renderers.get_renderer(RENDERER, tokenizer)
    rows = load_dataset("openai/gsm8k", "main", split="train").select(range(2))
    datums, spec = [], []
    for row in rows:
        convo = MathEnv.standard_fewshot_prefix() + [
            {"role": "user", "content": row["question"] + " Provide a numerical answer without units, written inside \\boxed{}."},
            {"role": "assistant", "content": row["answer"]}]
        mi, w = r.build_supervised_example(convo)
        toks = mi.to_ints(); w = [float(x) for x in w.tolist()]
        datums.append(types.Datum(model_input=types.ModelInput.from_ints(toks[:-1]),
                                  loss_fn_inputs={"target_tokens": toks[1:], "weights": w[1:]}))
        spec.append({"tokens": toks[:-1], "targets": toks[1:], "weights": w[1:]})
    return datums, spec, sha(jdump(spec))

def fingerprint(fwd):
    vals = []
    for out in fwd.loss_fn_outputs:
        vals.extend(float(x) for x in out["logprobs"].tolist())
    from research_exchange.restoration import output_fingerprint
    return vals, output_fingerprint(vals)

def adapter_hash(tc, rest, name):
    """save_weights_external -> download every file -> sha256 over safetensors tensors (name, shape, dtype, raw bytes)."""
    from safetensors.numpy import load as st_load
    path = tc.save_weights_external(name).result().path
    urls = rest.get_external_weights_urls(path).result().urls
    files, tensors_meta, h = {}, {}, hashlib.sha256()
    for fname in sorted(urls):
        data = urllib.request.urlopen(urls[fname], timeout=600).read()
        files[fname] = {"sha256": sha(data), "bytes": len(data)}
        if fname.endswith(".safetensors"):
            import safetensors, struct
            n = struct.unpack("<Q", data[:8])[0]; header = json.loads(data[8:8 + n]); base = 8 + n
            for k in sorted(x for x in header if x != "__metadata__"):
                m = header[k]; s, e = m["data_offsets"]; raw = data[base + s: base + e]
                h.update(jdump([k, m["shape"], m["dtype"]])); h.update(raw)
                import numpy as np
                is_zero = not any(raw)
                tensors_meta[k] = {"shape": m["shape"], "dtype": m["dtype"], "all_zero": is_zero}
    return {"external_path": path, "files": files, "tensors": tensors_meta, "adapter_sha256": h.hexdigest()}

def rank_summary(tensors):
    ranks = set(); b_zero = True; a_nonzero = True; n_a = n_b = 0
    for k, m in tensors.items():
        if "lora_A" in k:
            ranks.add(m["shape"][0]); n_a += 1; a_nonzero &= not m["all_zero"]
        if "lora_B" in k:
            ranks.add(m["shape"][1]); n_b += 1; b_zero &= m["all_zero"]
    return {"ranks_in_adapter": sorted(ranks), "n_lora_A": n_a, "n_lora_B": n_b, "all_B_zero": b_zero, "all_A_nonzero": a_nonzero}