#!/usr/bin/env python3
"""Merge the current best Actor LoRA adapter into MiniCPM5-1B."""

from __future__ import annotations

import argparse
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


DEFAULT_BASE_MODEL = "openbmb/MiniCPM5-1B"
DEFAULT_ADAPTER_DIR = Path("finetune/minicpm5-actor-lora")
DEFAULT_OUTPUT_DIR = Path("finetune/outputs/minicpm5-actor-merged")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base_model", default=DEFAULT_BASE_MODEL)
    parser.add_argument("--adapter_dir", type=Path, default=DEFAULT_ADAPTER_DIR)
    parser.add_argument("--output_dir", type=Path, default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
    parser.add_argument("--trust_remote_code", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    merge_lora(args)


def merge_lora(args: argparse.Namespace) -> None:
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not args.adapter_dir.exists():
        raise FileNotFoundError(f"Adapter directory does not exist: {args.adapter_dir}")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    dtype = resolve_torch_dtype(args.dtype, torch)
    print(f"Loading base model: {args.base_model}")
    base_model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        torch_dtype=dtype,
        device_map="auto",
        trust_remote_code=args.trust_remote_code,
    )
    print(f"Loading adapter: {args.adapter_dir}")
    model = PeftModel.from_pretrained(base_model, args.adapter_dir)
    print("Merging LoRA adapter into base model")
    merged_model = model.merge_and_unload()
    print(f"Saving merged model to: {args.output_dir}")
    merged_model.save_pretrained(str(args.output_dir), safe_serialization=True)

    tokenizer_source = args.adapter_dir if (args.adapter_dir / "tokenizer_config.json").exists() else args.base_model
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_source,
        trust_remote_code=args.trust_remote_code,
        use_fast=True,
    )
    tokenizer.save_pretrained(str(args.output_dir))
    write_merge_manifest(args.output_dir, args)
    print(f"Saved merged model to {args.output_dir}")


def resolve_torch_dtype(raw_dtype: str, torch: Any) -> Any:
    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[raw_dtype]


def write_merge_manifest(output_dir: Path, args: argparse.Namespace) -> None:
    manifest = {
        "base_model": args.base_model,
        "adapter_path": str(args.adapter_dir),
        "output_dir": str(args.output_dir),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "dtype": args.dtype,
        "trust_remote_code": args.trust_remote_code,
        "note": "v0 is the current best Actor LoRA candidate; v1 was useful for audit/eval tooling but did not outperform v0 on full eval.",
    }
    with (output_dir / "merge_manifest.json").open("w", encoding="utf-8") as handle:
        json.dump(manifest, handle, ensure_ascii=True, indent=2)
        handle.write("\n")


if __name__ == "__main__":
    main()