"""Publish a trained LoRA checkpoint to the Hub.

Usage:
    python export_adapter.py checkpoints/checkpoint-1000 you/LFM2.5-2.6B-opencode-lora-RL
    python export_adapter.py checkpoints/final you/LFM2.5-2.6B-opencode-lora-RL --also-merged
"""

import argparse
import shutil
import tempfile
from pathlib import Path

from huggingface_hub import HfApi

BASE_MODEL = "LiquidAI/LFM2.5-2.6B"

MODEL_CARD = """---
license: other
base_model: {base}
tags:
  - lora
  - rl
  - grpo
  - agentic-rl
  - opencode
---

# {repo} — LoRA adapter

LoRA adapter (rank {rank}) trained with async GRPO **inside the OpenCode agent harness**
on [SmolDataEnvs](https://huggingface.co/collections/FineEnvs/smoldataenvs) tasks, using
the OpenEnv capture proxy + Harbor + TRL stack from
[The ultimate guide to multi-harness RL](https://huggingface.co/spaces/FineEnvs/multi-harness-rl).

Fits `{base}` only.

## Load it

```python
from transformers import AutoModelForCausalLM
from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained("{base}")
model = PeftModel.from_pretrained(base, "{repo}")
```

```bash
# vLLM
vllm serve {base} --enable-lora --max-lora-rank {rank} \\
    --lora-modules opencode-lora={repo}
```
"""


def main():
    p = argparse.ArgumentParser()
    p.add_argument("checkpoint", type=Path, help="checkpoint dir containing adapter_model.safetensors")
    p.add_argument("repo_id", help="Hub repo for the adapter, e.g. you/LFM2.5-2.6B-opencode-lora-RL")
    p.add_argument("--also-merged", action="store_true",
                   help="also push a merged full model to <repo_id>-merged")
    args = p.parse_args()

    assert (args.checkpoint / "adapter_model.safetensors").exists(), \
        f"no adapter_model.safetensors in {args.checkpoint} — is this a PEFT checkpoint?"

    import json
    rank = json.loads((args.checkpoint / "adapter_config.json").read_text())["r"]

    api = HfApi()
    with tempfile.TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        shutil.copy(args.checkpoint / "adapter_model.safetensors", tmp)
        shutil.copy(args.checkpoint / "adapter_config.json", tmp)
        (tmp / "README.md").write_text(MODEL_CARD.format(base=BASE_MODEL, repo=args.repo_id, rank=rank))
        api.create_repo(args.repo_id, exist_ok=True)
        api.upload_folder(repo_id=args.repo_id, folder_path=tmp)
    print(f"adapter -> https://huggingface.co/{args.repo_id}")

    if args.also_merged:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from peft import PeftModel

        base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
        model = PeftModel.from_pretrained(base, args.checkpoint)
        merged = model.merge_and_unload()
        merged_id = f"{args.repo_id}-merged"
        api.create_repo(merged_id, exist_ok=True)
        merged.push_to_hub(merged_id)
        AutoTokenizer.from_pretrained(BASE_MODEL).push_to_hub(merged_id)
        print(f"merged  -> https://huggingface.co/{merged_id}")


if __name__ == "__main__":
    main()