"""
Merge LoRA adapter weights into the base model and save a full checkpoint.

Usage:
    python finetune/merge_lora.py
    python finetune/merge_lora.py --lora_path finetune/checkpoints --output merged_model
"""

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT / "src"))

from model_utils import load_config  # noqa: E402

try:
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer
except ImportError as e:
    print(f"Missing dependency: {e}")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--lora_path", default=None)
    parser.add_argument("--output", default=None)
    parser.add_argument("--base_model", default=None,
                        help="Override base model path (e.g. merged_model)")
    parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
    args = parser.parse_args()

    cfg = load_config()
    if args.base_model:
        model_id = str(ROOT / args.base_model)
    else:
        model_id = cfg["model"]["local_path"] or cfg["model"]["id"]
    lora_path = Path(args.lora_path or ROOT / cfg["finetune"]["output_dir"])
    output_path = Path(args.output or ROOT / "merged_model")

    if not lora_path.exists():
        print(f"LoRA checkpoint not found: {lora_path}")
        sys.exit(1)

    dtype = getattr(torch, args.dtype)
    # Force CPU to avoid VRAM OOM during merge; RAM usage ~18GB for 9B bf16
    print(f"Loading base model on CPU in {args.dtype} (may take a few minutes)...")
    base_model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map="cpu",
        low_cpu_mem_usage=False,  # force full materialisation, no meta tensors
        trust_remote_code=True,
    )

    print(f"Loading LoRA adapter from {lora_path}...")
    model = PeftModel.from_pretrained(base_model, str(lora_path))

    print("Merging weights...")
    model = model.merge_and_unload()

    output_path.mkdir(parents=True, exist_ok=True)
    print(f"Saving merged model to {output_path} (sharded 4GB per file)...")
    model.save_pretrained(
        str(output_path),
        safe_serialization=True,
        max_shard_size="4GB",
    )

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    # model_id may be a local path; always copy tokenizer alongside merged weights
    tokenizer.save_pretrained(str(output_path))

    print(f"Done. Merged model saved to {output_path}")


if __name__ == "__main__":
    main()