#!/usr/bin/env python3
"""fastloop_merge_adapter.py <base-dir> <adapter-dir> <merged-dir> — re-create the
BF16 merged checkpoint from a kept adapter (the ladder runners delete merged/
after eval; when an eval fails for plumbing reasons this rebuilds it in ~5 min
instead of a 2 h retune). Same merge law as fastloop_smoke_train.py M gate."""
import os, shutil, sys, torch
from peft import PeftModel
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoTokenizer
base, adapter, merged_dir = sys.argv[1:4]
cfg = AutoConfig.from_pretrained(base)
cls = AutoModelForImageTextToText if hasattr(cfg, "text_config") else AutoModelForCausalLM
m = cls.from_pretrained(base, device_map="auto", dtype=torch.bfloat16)
m = PeftModel.from_pretrained(m, adapter).merge_and_unload()
if os.path.isdir(merged_dir): shutil.rmtree(merged_dir)
m.save_pretrained(merged_dir, safe_serialization=True)
AutoTokenizer.from_pretrained(base).save_pretrained(merged_dir)
for f in os.listdir(base):
    if f.endswith(".index.json"): continue
    if f.endswith((".jinja", ".json", ".py", ".txt", ".model")) and not os.path.exists(os.path.join(merged_dir, f)):
        shutil.copy(os.path.join(base, f), merged_dir)
st = [f for f in os.listdir(merged_dir) if f.endswith(".safetensors")]
print(f"MERGE PASS: {len(st)} shards -> {merged_dir}" if st else "MERGE FAIL")
sys.exit(0 if st else 1)