#!/usr/bin/env python3
"""Merge a selected LoRA adapter into the base model."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import torch
import yaml


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser.add_argument("--adapter", help="Adapter checkpoint path. Defaults to run.output_dir.")
    parser.add_argument("--output", help="Merged output path. Defaults to run.merged_output_dir.")
    parser.add_argument("--dtype", default=None, choices=["bfloat16", "float16", "float32"])
    return parser.parse_args()


def read_yaml(path: str | Path) -> dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as fh:
        payload = yaml.safe_load(fh) or {}
    if not isinstance(payload, dict):
        raise TypeError(f"Expected a YAML mapping in {path}")
    return payload


def torch_dtype(name: str):
    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[name]


def load_base_model(transformers_module, model_name: str, dtype, trust_remote_code: bool):
    errors: list[str] = []
    for class_name in [
        "AutoModelForMultimodalLM",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForCausalLM",
    ]:
        model_cls = getattr(transformers_module, class_name, None)
        if model_cls is None:
            continue
        try:
            return model_cls.from_pretrained(
                model_name,
                torch_dtype=dtype,
                trust_remote_code=trust_remote_code,
                device_map="auto",
            )
        except Exception as exc:
            errors.append(f"{class_name}: {exc!r}")
    raise RuntimeError("Could not load base model:\n" + "\n".join(errors))


def main() -> int:
    args = parse_args()
    config = read_yaml(args.config)
    run_cfg = config["run"]
    model_cfg = config["model"]

    adapter_path = Path(args.adapter or run_cfg["output_dir"])
    output_path = Path(args.output or run_cfg["merged_output_dir"])
    dtype_name = args.dtype or str(config.get("merge", {}).get("merged_dtype", model_cfg.get("dtype", "bfloat16")))

    import transformers
    from peft import PeftModel
    from transformers import AutoTokenizer

    dtype = torch_dtype(dtype_name)
    base = load_base_model(
        transformers,
        model_cfg["name_or_path"],
        dtype=dtype,
        trust_remote_code=bool(model_cfg.get("trust_remote_code", True)),
    )
    model = PeftModel.from_pretrained(base, adapter_path)
    merged = model.merge_and_unload()
    output_path.mkdir(parents=True, exist_ok=True)
    merged.save_pretrained(output_path, safe_serialization=True)

    tokenizer = AutoTokenizer.from_pretrained(model_cfg["name_or_path"], trust_remote_code=True)
    tokenizer.save_pretrained(output_path)

    manifest = {
        "base_model": model_cfg["name_or_path"],
        "adapter": str(adapter_path),
        "output": str(output_path),
        "dtype": dtype_name,
    }
    (output_path / "merge_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(manifest, indent=2))
    return 0


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