"""Materialize one PEFT checkpoint for single-model evaluation.

The merged directory is an evaluation cache, not a training checkpoint.  It
contains no optimizer state and may be deleted after its generations have been
durably written.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-model", required=True)
    parser.add_argument("--adapter", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    import torch
    from peft import PeftModel
    from transformers import (
        AutoConfig,
        AutoModelForCausalLM,
        AutoModelForImageTextToText,
        AutoTokenizer,
    )

    base = Path(args.base_model).resolve()
    adapter = Path(args.adapter).resolve()
    output = Path(args.output).resolve()
    if not base.is_dir() or not adapter.is_dir():
        raise SystemExit("--base-model and --adapter must be existing local directories")
    if output.exists() and any(output.iterdir()):
        raise SystemExit(f"refusing to overwrite non-empty merged directory: {output}")
    output.mkdir(parents=True, exist_ok=True)

    adapter_weights = adapter / "adapter_model.safetensors"
    if not adapter_weights.is_file():
        raise SystemExit(f"adapter weights not found: {adapter_weights}")

    tokenizer = AutoTokenizer.from_pretrained(base, trust_remote_code=True, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        base,
        trust_remote_code=True,
        local_files_only=True,
        dtype=torch.bfloat16,
        device_map="cpu",
    )
    merged = PeftModel.from_pretrained(model, adapter, local_files_only=True).merge_and_unload()
    base_config = AutoConfig.from_pretrained(base, trust_remote_code=True, local_files_only=True)
    if base_config.model_type == "qwen3_5" and hasattr(base_config, "vision_config"):
        # AutoModelForCausalLM intentionally extracts Qwen3.5's text model.  A
        # bare Qwen3_5TextConfig cannot currently be loaded by vLLM 0.22.x,
        # which still expects the outer config's vision_config even in
        # language-model-only mode.  Put the merged text tower and LM head back
        # into the original full-model container so direct evaluation remains
        # loadable while still using the merged adapter weights.
        full_model = AutoModelForImageTextToText.from_pretrained(
            base,
            trust_remote_code=True,
            local_files_only=True,
            dtype=torch.bfloat16,
            device_map="cpu",
        )
        full_model.model.language_model = merged.model
        full_model.lm_head = merged.lm_head
        model_to_save = full_model
        container = "full_qwen3_5"
    else:
        model_to_save = merged
        container = "causal_lm"
    model_to_save.save_pretrained(output, safe_serialization=True, max_shard_size="5GB")
    tokenizer.save_pretrained(output)
    manifest = {
        "schema": "merged_peft_eval_cache_v1",
        "base_model": str(base),
        "adapter": str(adapter),
        "adapter_sha256": _sha256(adapter_weights),
        "dtype": "bfloat16",
        "container": container,
        "optimizer_state_included": False,
    }
    temporary = output / "merge_manifest.json.tmp"
    temporary.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    with temporary.open("r+") as handle:
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, output / "merge_manifest.json")
    print(f"[merge] {adapter} + {base} -> {output}")


if __name__ == "__main__":
    main()