import os, torch, random
from contextlib import contextmanager

def set_seed(seed):
    random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

@contextmanager
def autocast_if_bf16(enabled=True):
    if enabled and torch.cuda.is_available() and torch.cuda.is_bf16_supported():
        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
            yield
    else:
        yield

def save_checkpoint(model, opt, step, ckpt_dir):
    os.makedirs(ckpt_dir, exist_ok=True)
    path = os.path.join(ckpt_dir, f"oai_flash_step{step}.pt")
    torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, path)
    print("Saved:", path)

def latest_checkpoint(ckpt_dir):
    if not os.path.exists(ckpt_dir): return None
    files = [f for f in os.listdir(ckpt_dir) if f.startswith("oai_flash_step")]
    if not files: return None
    files.sort(key=lambda x: int(x.split("step")[-1].split(".")[0]))
    return os.path.join(ckpt_dir, files[-1])