"""Merge local LoRA adapters into an explicitly supplied, fingerprinted base."""

from __future__ import annotations

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

from .evaluation import checkpoint_sha256
from .training import read_json, require_versions, sha256, write_manifest


def merge_adapter(
    base_dir: str | Path,
    adapter_dir: str | Path,
    output_dir: str | Path,
    *,
    dtype: str = "bfloat16",
    device: str = "cpu",
) -> dict[str, Any]:
    """Merge with PEFT, using no Hub lookup and ignoring the adapter's base pointer.

    Output is a standalone full checkpoint with the base tokenizer. An adjacent
    ``<output>.merge-manifest.json`` binds base, adapter and output tree hashes.
    The adjacent manifest keeps the checkpoint hash free of self-reference.
    """
    if dtype not in {"bfloat16", "float32", "float16"}:
        raise ValueError("dtype must be bfloat16, float32 or float16")
    if not isinstance(device, str) or re.fullmatch(r"cpu|cuda(?::[0-9]+)?", device) is None:
        raise ValueError("device must be cpu, cuda or cuda:N")
    base, adapter, output = (Path(p).resolve() for p in (base_dir, adapter_dir, output_dir))
    manifest_path = output.with_name(output.name + ".merge-manifest.json")
    if output.exists() or manifest_path.exists():
        raise ValueError("Merged checkpoint output or manifest already exists")
    for label, source in (("base", base), ("adapter", adapter)):
        if not source.is_dir():
            raise ValueError(f"{label} must be an existing local directory")
        if output == source or output.is_relative_to(source) or source.is_relative_to(output):
            raise ValueError("Merged output must be separate from source directories")
    if not (base / "config.json").is_file() or (base / "adapter_config.json").exists():
        raise ValueError("base must be a full model checkpoint, not another adapter")
    if not (base / "tokenizer_config.json").is_file():
        raise ValueError("base must include its tokenizer configuration")
    if not (adapter / "adapter_config.json").is_file() or not (adapter / "adapter_model.safetensors").is_file():
        raise ValueError("adapter requires adapter_config.json and adapter_model.safetensors")
    adapter_config = read_json(adapter / "adapter_config.json")
    if adapter_config.get("peft_type") != "LORA":
        raise ValueError("Only LoRA adapter checkpoints are supported")
    base_hash, adapter_hash = checkpoint_sha256(base), checkpoint_sha256(adapter)
    installed = require_versions(("torch", "transformers", "peft"))
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if device.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but is unavailable; use the GPU allocation")
    model = AutoModelForCausalLM.from_pretrained(
        str(base), dtype=getattr(torch, dtype), local_files_only=True,
        trust_remote_code=False, use_safetensors=True,
    )
    model.to(device)
    tokenizer = AutoTokenizer.from_pretrained(str(base), local_files_only=True, trust_remote_code=False)
    # The explicitly instantiated base object determines the actual base weights;
    # adapter_config.base_model_name_or_path is provenance only, never followed.
    adapted = PeftModel.from_pretrained(model, str(adapter), local_files_only=True, is_trainable=False)
    if checkpoint_sha256(base) != base_hash or checkpoint_sha256(adapter) != adapter_hash:
        raise ValueError("Base or adapter changed while loading; no merged checkpoint was written")
    merged = adapted.merge_and_unload(safe_merge=True)
    output.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix=".merge-", dir=output.parent) as temporary:
        staging = Path(temporary) / "checkpoint"
        merged.save_pretrained(str(staging), safe_serialization=True)
        tokenizer.save_pretrained(str(staging))
        if (staging / "adapter_config.json").exists():
            raise ValueError("Merge did not produce a standalone full checkpoint")
        if checkpoint_sha256(base) != base_hash or checkpoint_sha256(adapter) != adapter_hash:
            raise ValueError("Base or adapter changed during merging; no checkpoint was published")
        output_hash = checkpoint_sha256(staging)
        tokenizer_hashes = {
            file.name: sha256(file) for file in sorted(staging.iterdir())
            if file.is_file() and ("token" in file.name or file.name in {
                "chat_template.jinja", "vocab.json", "merges.txt", "special_tokens_map.json",
                "added_tokens.json", "sentencepiece.bpe.model", "spiece.model", "tokenizer.model",
            })
        }
        report = {
            "schema_version": 1, "status": "completed", "method": "peft_lora_merge",
            "base_model": {"path": str(base), "sha256": base_hash},
            "adapter": {"path": str(adapter), "sha256": adapter_hash,
                        "declared_base_model_name_or_path": adapter_config.get("base_model_name_or_path"),
                        "base_pointer_used": False},
            "output": {"path": str(output), "sha256": output_hash},
            "tokenizer_source": str(base), "tokenizer_file_sha256": tokenizer_hashes,
            "dtype": dtype, "device": device, "dependencies": installed,
            "completed_at": datetime.now(timezone.utc).isoformat(),
            "manifest": str(manifest_path), "published": False,
        }
        # Keep a complete staging checkpoint until the rename succeeds.
        if output.exists() or manifest_path.exists():
            raise ValueError("Merged checkpoint output appeared while merging")
        staging.rename(output)
        write_manifest(manifest_path, report)
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("base", type=Path)
    parser.add_argument("adapter", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--dtype", default="bfloat16", choices=("bfloat16", "float32", "float16"))
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()
    try:
        print(json.dumps(merge_adapter(args.base, args.adapter, args.output,
                                       dtype=args.dtype, device=args.device), indent=2))
    except (ValueError, OSError, RuntimeError, ImportError) as error:
        parser.exit(2, f"error: {error}\n")


if __name__ == "__main__":
    main()