"""Merge a LoRA adapter (+ trained readout) into a d3 package, keeping the package layout.

    python -m d3train.merge --base <d3 package> --adapter <run>/final --out <merged package> [--name d3-mini-ft]

The base shards are rewritten tensor by tensor under the same names, shard files and index: an adapted weight
becomes ``(W.float() + scale * B @ A).to(W.dtype)``, every other tensor is copied unchanged (so a zero-step adapter,
B = 0, reproduces every tensor bit for bit). ``readout.safetensors`` comes from the run (FP32, [255, H]);
config.json, decision_config.json (temperature optionally replaced), tokenizer, chat template and the runtime code
are copied. MODEL_MANIFEST.json is regenerated for the new files, so ``verify="fast"`` keeps working.
"""

from __future__ import annotations

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

import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file

WEIGHTS = re.compile(r"^model(-\d{5}-of-\d{5})?\.safetensors$")
REGENERATED = {"MODEL_MANIFEST.json", "readout.safetensors", "model.safetensors.index.json", "decision_config.json"}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        for block in iter(lambda: stream.read(16 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def lora_deltas(adapter: Path) -> dict[str, torch.Tensor]:
    """{checkpoint weight name: FP32 delta} from a PEFT LoRA adapter saved by d3train.train."""
    config = json.loads((adapter / "adapter_config.json").read_text(encoding="utf-8"))
    if config.get("peft_type") != "LORA" or config.get("use_dora") or config.get("fan_in_fan_out"):
        raise ValueError("only plain LoRA adapters (no DoRA, fan_in_fan_out False) are supported")
    r, alpha = config["r"], config["lora_alpha"]
    if config.get("rank_pattern") or config.get("alpha_pattern"):
        raise ValueError("rank/alpha patterns are not supported")
    scale = alpha / (r ** 0.5) if config.get("use_rslora") else alpha / r
    tensors = load_file(str(adapter / "adapter_model.safetensors"))
    a_suffix, b_suffix = ".lora_A.weight", ".lora_B.weight"
    deltas = {}
    for name, a in tensors.items():
        if not name.endswith(a_suffix):
            continue
        module = name[: -len(a_suffix)]
        b = tensors[module + b_suffix]
        key = module.removeprefix("base_model.model.") + ".weight"
        deltas[key] = (b.float() @ a.float()) * scale
    unknown = [n for n in tensors if not n.endswith((a_suffix, b_suffix))]
    if unknown:
        raise ValueError(f"adapter holds tensors other than LoRA A/B: {unknown[:4]}")
    return deltas


def check_llamacpp_acceptance(root: Path) -> list[str]:
    """What llama.cpp's conversion/vllm_d3.py + pplx_decider.py check before converting; [] when accepted."""
    problems = []
    for name in ("config.json", "decision_config.json", "readout.safetensors"):
        if not (root / name).is_file():
            problems.append(f"missing {name}")
    if problems:
        return problems
    decision = json.loads((root / "decision_config.json").read_text(encoding="utf-8"))
    if decision.get("format_id") != "d3-code-readout-v1":
        problems.append("format_id is not d3-code-readout-v1 (not detected as vllm-d3)")
    if decision.get("prompt") != "d3":
        problems.append("prompt != d3")
    if decision.get("attention_mode") != "noncausal_full_attention":
        problems.append("attention_mode != noncausal_full_attention")
    if decision.get("pooling") != "last":
        problems.append("pooling != last")
    if not isinstance(decision.get("temperature"), (int, float)):
        problems.append("temperature missing (written for choice/score/noul)")
    with safe_open(str(root / "readout.safetensors"), framework="pt") as handle:
        if list(handle.keys()) != ["weight"]:
            problems.append("readout.safetensors must hold exactly 'weight'")
        else:
            shape = handle.get_slice("weight").get_shape()
            ids = decision.get("token_ids") or []
            if not (shape[0] == len(ids) == len(set(ids))):
                problems.append("readout rows != number of distinct token_ids")
    config = json.loads((root / "config.json").read_text(encoding="utf-8"))
    vocab = (config.get("text_config") or config).get("vocab_size")
    if any(not 0 <= i < (vocab or 0) for i in decision.get("token_ids") or []):
        problems.append("token id outside vocab")
    if config.get("model_type") != "qwen3_5":
        problems.append("model_type != qwen3_5 (converter maps vllm-d3 to MODEL_ARCH.QWEN35)")
    return problems


def merge(
    base: Path, adapter: Path, out: Path, *, name: str | None = None, temperature: float | None = None,
) -> dict[str, Any]:
    base, adapter, out = Path(base), Path(adapter), Path(out)
    if out.exists() and any(out.iterdir()):
        raise FileExistsError(f"{out} is not empty")
    out.mkdir(parents=True, exist_ok=True)
    deltas = lora_deltas(adapter) if (adapter / "adapter_config.json").is_file() else {}
    readout_file = adapter / "readout.safetensors"
    readout = load_file(str(readout_file if readout_file.is_file() else base / "readout.safetensors"))["weight"]
    if readout.dtype != torch.float32:
        raise ValueError("the readout must be FP32")

    # copy everything that is not weights or regenerated metadata
    for path in sorted(base.rglob("*")):
        rel = path.relative_to(base)
        if path.is_dir() or any(p.startswith(".") or p == "__pycache__" for p in rel.parts):
            continue
        if rel.as_posix() in REGENERATED or WEIGHTS.match(rel.name) or rel.name.endswith(".header.json"):
            continue
        target = out / rel
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target)

    # weights: same names, same shards
    index_path = base / "model.safetensors.index.json"
    if index_path.is_file():
        index = json.loads(index_path.read_text(encoding="utf-8"))
        shards = sorted(set(index["weight_map"].values()))
    else:
        index, shards = None, ["model.safetensors"]
    applied, changed = set(), 0
    for shard in shards:
        tensors, metadata = {}, None
        with safe_open(str(base / shard), framework="pt") as handle:
            metadata = handle.metadata()
            for key in handle.keys():
                tensor = handle.get_tensor(key)
                if key in deltas:
                    delta = deltas[key]
                    if delta.shape != tensor.shape:
                        raise ValueError(f"{key}: delta {tuple(delta.shape)} vs weight {tuple(tensor.shape)}")
                    merged = (tensor.float() + delta).to(tensor.dtype)
                    changed += int(not torch.equal(merged, tensor))
                    tensor = merged
                    applied.add(key)
                tensors[key] = tensor.contiguous()
        save_file(tensors, str(out / shard), metadata=metadata)
    missing = sorted(set(deltas) - applied)
    if missing:
        raise ValueError(f"adapter weights not found in the base checkpoint: {missing[:4]}")
    if index is not None:
        # same tensor names, shapes and dtypes: the index (and its total_size) is unchanged
        shutil.copy2(index_path, out / "model.safetensors.index.json")

    save_file({"weight": readout.contiguous()}, str(out / "readout.safetensors"))
    decision = json.loads((base / "decision_config.json").read_text(encoding="utf-8"))
    if temperature is not None:
        decision["temperature"] = float(temperature)
    (out / "decision_config.json").write_text(json.dumps(decision, indent=2) + "\n", encoding="utf-8")
    notice = out / "NOTICE"
    if notice.is_file():
        notice.write_text(
            notice.read_text(encoding="utf-8").rstrip("\n")
            + f"\n\nModified: LoRA fine-tune merged into the text backbone and readout ({datetime.now(timezone.utc).date()}).\n",
            encoding="utf-8",
        )

    problems = check_llamacpp_acceptance(out)
    report = {"adapted_tensors": len(applied), "changed_tensors": changed, "llamacpp_acceptance": problems or "ok"}
    (out / "MERGE-REPORT.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")

    # manifest for the runtime's verify="fast"
    old = {}
    if (base / "MODEL_MANIFEST.json").is_file():
        old = json.loads((base / "MODEL_MANIFEST.json").read_text(encoding="utf-8"))
    files = sorted(
        p.relative_to(out).as_posix() for p in out.rglob("*")
        if p.is_file() and p.name != "MODEL_MANIFEST.json"
        and not any(part.startswith(".") or part == "__pycache__" for part in p.relative_to(out).parts)
    )
    manifest = {
        "schema": "d3-package-manifest/1",
        **{k: old[k] for k in ("format_id", "prompt", "attention_mode", "max_input_tokens", "base_model", "backbone",
                               "licence", "parameters") if k in old},
        "model_name": name or f"{old.get('model_name', base.name)}-ft",
        "derived_from": {"model_name": old.get("model_name"), "repo_id": old.get("repo_id"),
                         "built_utc": old.get("built_utc"), "adapter": adapter.name},
        "built_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "files_sha256": {f: sha256_file(out / f) for f in files},
        "files_bytes": {f: (out / f).stat().st_size for f in files},
    }
    (out / "MODEL_MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    return report


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--base", required=True, type=Path)
    ap.add_argument("--adapter", required=True, type=Path, help="run checkpoint dir (adapter + readout.safetensors)")
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--name")
    ap.add_argument("--temperature", type=float, help="write a recalibrated temperature into decision_config.json")
    args = ap.parse_args()
    print(json.dumps(merge(args.base, args.adapter, args.out, name=args.name, temperature=args.temperature)))


if __name__ == "__main__":
    main()