import asyncio, hashlib, json, os, time
import numpy as np
import tinker
from tinker_cookbook import renderers
from tinker_cookbook.recipes.math_rl.math_env import get_tokenizer, MathEnv

RUN = "run_cbb327f6ab"
MODEL = "Qwen/Qwen3.5-2B"
RENDERER = "qwen3_5_disable_thinking"
LORA_RANK = 32
LR = 4e-5
ADAM = tinker.AdamParams(learning_rate=LR, beta1=0.9, beta2=0.95, eps=1e-8)
MAX_TOKENS = 512
TEMPERATURE = 1.0
TAU = 0.1
ARMS = {"grpo_gs8_seed0": dict(group_size=8, groups=8), "gs1_seed0": dict(group_size=1, groups=64)}
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "out")
os.makedirs(OUT, exist_ok=True)


def service_client():
    return tinker.ServiceClient(
        base_url=os.environ["TINKER_BASE_URL"], api_key=os.environ["TINKER_API_KEY"],
        default_headers_overrides={"Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"],
                                   "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"]})


def renderer_and_tokenizer():
    tok = get_tokenizer(MODEL)
    return renderers.get_renderer(RENDERER, tokenizer=tok), tok


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


def sha256_json(o):
    return sha256_bytes(json.dumps(o, sort_keys=True, default=str).encode())


def dump(name, obj):
    p = os.path.join(OUT, name)
    with open(p, "w") as f:
        json.dump(obj, f, indent=1, default=str)
    return p


async def with_capacity_retry(fn, budget_s=600, log=print):
    t0, delay = time.time(), 15
    while True:
        try:
            return await fn()
        except Exception as e:
            msg = str(e)
            if time.time() - t0 > budget_s or not any(k in msg.lower() for k in ("capacity", "429", "503", "unavailable", "exhaust")):
                raise
            log(f"capacity retry after error: {msg[:300]}")
            await asyncio.sleep(delay)
            delay = min(delay * 2, 120)