# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "unsloth",
#     "transformers>=4.57",
#     "huggingface_hub",
# ]
# ///
"""Merge the SecureCoder LoRA into the base Qwen3-Coder-30B-A3B-Instruct
checkpoint, upload a 16-bit safetensors repo, then quantise to Q4_K_M GGUF.

Default base: unsloth/Qwen3-Coder-30B-A3B-Instruct
Default adapter: Taimwe/securecoder-30b-pro

Run on HF Jobs (a100-large has the headroom to load 30B in 16-bit):
    hf jobs run -d --flavor a100-large --timeout 90m --secrets HF_TOKEN \\
      ghcr.io/astral-sh/uv:python3.12-bookworm \\
      uv run --no-project https://huggingface.co/Taimwe/securecoder-scripts/resolve/main/merge_securecoder.py \\
      -- --adapter Taimwe/securecoder-30b-pro \\
         --output-repo Taimwe/securecoder-30b-pro-merged \\
         --gguf-repo Taimwe/securecoder-30b-pro-GGUF
"""

from __future__ import annotations

import argparse
import logging
import os
import shutil
import sys
import time

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("merge")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Merge + GGUF + push the SecureCoder LoRA")
    p.add_argument("--base", default="unsloth/Qwen3-Coder-30B-A3B-Instruct")
    p.add_argument("--adapter", default="Taimwe/securecoder-30b-pro")
    p.add_argument("--output-repo", default="Taimwe/securecoder-30b-pro-merged")
    p.add_argument("--gguf-repo", default=None)
    p.add_argument("--gguf-quant", default="Q4_K_M")
    p.add_argument("--private", action="store_true")
    p.add_argument("--work-dir", default="/data/securecoder-merge")
    p.add_argument("--max-shard-size", default="5GB")
    return p.parse_args()

def main() -> int:
    args = parse_args()
    token = os.environ.get("HF_TOKEN")
    if not token:
        log.error("HF_TOKEN not set")
        return 1

    import torch
    from huggingface_hub import HfApi
    from unsloth import FastLanguageModel

    if not torch.cuda.is_available():
        log.error("no CUDA - merge needs a GPU")
        return 1
    log.info("GPU: %s", torch.cuda.get_device_name(0))

    work = args.work_dir
    if os.path.exists(work):
        shutil.rmtree(work)
    os.makedirs(work, exist_ok=True)

    log.info("loading base %s in 16-bit ...", args.base)
    started = time.time()
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.base,
        max_seq_length=2048,
        dtype=torch.bfloat16,
        load_in_4bit=False,
    )

    log.info("loading adapter %s ...", args.adapter)
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, args.adapter, token=token)
    log.info("merging ...")
    model = model.merge_and_unload()
    log.info("merge done in %.1f min", (time.time() - started) / 60)

    out_dir = os.path.join(work, "merged")
    model.save_pretrained(out_dir, safe_serialization=True, max_shard_size=args.max_shard_size)
    tokenizer.save_pretrained(out_dir)
    log.info("saved merged model to %s", out_dir)

    api = HfApi(token=token)
    api.create_repo(args.output_repo, repo_type="model", exist_ok=True, private=args.private)
    log.info("uploading to %s ...", args.output_repo)
    api.upload_folder(folder_path=out_dir, repo_id=args.output_repo, repo_type="model",
                      commit_message="Merge SecureCoder LoRA into base (16-bit)")
    log.info("merged model live: https://huggingface.co/%s", args.output_repo)

    if args.gguf_repo:
        log.info("re-loading merged model for GGUF export ...")
        from unsloth import FastLanguageModel as FLM
        model, tokenizer = FLM.from_pretrained(
            model_name=out_dir,
            max_seq_length=2048,
            dtype=torch.bfloat16,
            load_in_4bit=False,
        )
        gguf_path = os.path.join(work, "gguf")
        os.makedirs(gguf_path, exist_ok=True)
        log.info("quantising to %s ...", args.gguf_quant)
        try:
            model.quantize_gguf_model(save_dir=gguf_path, quantization=args.gguf_quant)
        except Exception as exc:  # noqa: BLE001
            log.warning("model.quantize_gguf_model failed (%s); falling back to llama-quantize CLI", exc)
            from huggingface_hub import hf_hub_download
            from pathlib import Path as _P
            qbin = hf_hub_download(repo_id="unsloth/llama.cpp", filename="llama-quantize",
                                   repo_type="model", token=token)
            import subprocess
            subprocess.run(["chmod", "+x", qbin], check=False)
            src = next(_P(out_dir).glob("*.gguf"), None)
            if src is None:
                log.error("no GGUF produced by Unsloth quantise pass")
                return 1
            subprocess.run([qbin, str(src), str(_P(gguf_path) / src.name), args.gguf_quant], check=True)

        api.create_repo(args.gguf_repo, repo_type="model", exist_ok=True, private=args.private)
        api.upload_folder(folder_path=gguf_path, repo_id=args.gguf_repo, repo_type="model",
                          commit_message=f"Add {args.gguf_quant} GGUF export")
        log.info("GGUF live: https://huggingface.co/%s", args.gguf_repo)

    print("=" * 78)
    print("MERGE COMPLETE")
    print(f"  merged: https://huggingface.co/{args.output_repo}")
    if args.gguf_repo:
        print(f"  gguf  : https://huggingface.co/{args.gguf_repo}")
    print("=" * 78)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())