import os, json, hashlib, sys, time
sys.path.insert(0, os.path.expanduser("~/repos/research-exchange"))
import tinker

RUN = "run_b89b93886b"
MODEL = "Qwen/Qwen2.5-3B"
RENDERER = "role_colon"
HEADERS = {"Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"], "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"]}


def service_client(**kw):
    return tinker.ServiceClient(base_url=os.environ["TINKER_BASE_URL"], default_headers_overrides=HEADERS, **kw)


def sha(obj):
    return hashlib.sha256(json.dumps(obj, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()).hexdigest()


def fixed_datum():
    """One fixed GSM8K train example (row 0) rendered with the recipe prompt + reference solution."""
    from datasets import load_dataset
    from tinker_cookbook import renderers
    from tinker_cookbook.tokenizer_utils import get_tokenizer
    from tinker_cookbook.recipes.math_rl.math_env import MathEnv, Gsm8kDataset, extract_gsm8k_final_answer
    tok = get_tokenizer(MODEL)
    r = renderers.get_renderer(RENDERER, tokenizer=tok)
    row = load_dataset("openai/gsm8k", "main", split="train")[0]
    convo = MathEnv.standard_fewshot_prefix() + [{"role": "user", "content": row["question"] + Gsm8kDataset.question_suffix()}]
    prompt = r.build_generation_prompt(convo).to_ints()
    sol = row["answer"].split("####")[0].strip()
    comp = tok.encode(" " + sol + " \\boxed{" + extract_gsm8k_final_answer(row["answer"]) + "}", add_special_tokens=False)
    toks = prompt + comp
    targets = toks[1:]
    weights = [0.0] * (len(prompt) - 1) + [1.0] * len(comp)
    batch = {"tokens": toks[:-1], "targets": targets, "weights": weights}
    datum = tinker.Datum(model_input=tinker.ModelInput.from_ints(batch["tokens"]),
                         loss_fn_inputs={"target_tokens": tinker.TensorData(data=targets, dtype="int64", shape=[len(targets)]),
                                         "weights": tinker.TensorData(data=weights, dtype="float32", shape=[len(weights)])})
    return datum, batch


def forward_fingerprint(tc, datum):
    from research_exchange.restoration import output_fingerprint
    out = tc.forward([datum], "cross_entropy").result()
    lp = [float(x) for x in out.loss_fn_outputs[0]["logprobs"].tolist()]
    return lp, output_fingerprint(lp)


def info_dump(tc):
    i = tc.get_info()
    d = i.model_dump(mode="json")
    return d