#!/usr/bin/env python3
"""Merge one tensor at a time, then export GGUF without loading a full 7B CPU model."""
from __future__ import annotations

import argparse
import gc
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

BASE = "Qwen/Qwen2.5-Coder-7B-Instruct"
BASE_REVISION = "c03e6d358207e414f1eca0bb1891e29f1db0e242"
LLAMA_REVISION = "11fe02151f79c41d0d4af7da708755d73b9c0da6"


def sha256(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for b in iter(lambda: f.read(1024 * 1024), b""):
            h.update(b)
    return h.hexdigest()


def atomic_json(path, data):
    tmp = path.with_name(path.name + ".tmp")
    with tmp.open("w") as f:
        json.dump(data, f, indent=2)
        f.flush()
        os.fsync(f.fileno())
    tmp.replace(path)


def lora_pairs(adapter_file):
    from safetensors import safe_open
    pairs = {}
    with safe_open(adapter_file, framework="pt", device="cpu") as f:
        for name in f.keys():  # noqa: SIM118 - safetensors reader is not a dict
            match = re.fullmatch(r"(?:base_model\.model\.)?(.+)\.lora_([AB])(?:\.default)?\.weight", name)
            if not match:
                raise ValueError(f"Unsupported adapter tensor: {name}; only plain LoRA A/B supported")
            pairs.setdefault(match[1] + ".weight", {})[match[2]] = f.get_tensor(name)
    if not pairs or any(set(pair) != {"A", "B"} for pair in pairs.values()):
        raise ValueError("Adapter contains missing LoRA A/B tensor pairs")
    return pairs


def merge_tensor(weight, pair, scale, row_chunk=128):
    """Keep full output in fp16; calculate float32 deltas in small row blocks."""
    import torch
    # A dtype conversion already allocates. Chaining clone() would transiently
    # duplicate the 1GiB vocabulary tensor on the 12GiB free Colab host.
    result = weight.to(dtype=torch.float16, copy=True)
    if pair is None:
        return result
    if weight.ndim != 2 or pair["A"].shape[1] != weight.shape[1] or pair["B"].shape[0] != weight.shape[0]:
        raise ValueError("LoRA matrix dimensions do not match original base weights")
    a = pair["A"].float()
    for start in range(0, weight.shape[0], row_chunk):
        stop = min(start + row_chunk, weight.shape[0])
        # Read original precision here instead of the rounded fp16 result.
        block = weight[start:stop].float()
        block.addmm_(pair["B"][start:stop].float(), a, beta=1, alpha=scale)
        result[start:stop] = block.to(torch.float16)
    if not torch.isfinite(result).all():
        raise ValueError("Nonfinite weights after fp16 merge")
    return result.contiguous()


def merge(adapter, output, cache):
    import torch
    from huggingface_hub import snapshot_download
    from safetensors import safe_open
    from safetensors.torch import save_file
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
    config = json.loads((adapter / "adapter_config.json").read_text())
    if config.get("use_dora") or config.get("use_rslora") or config.get("fan_in_fan_out") or config.get("modules_to_save") or config.get("rank_pattern") or config.get("alpha_pattern") or config.get("bias", "none") != "none":
        raise ValueError("Export supports only the fixed plain LoRA configuration used by train.py")
    if config.get("r") != 32 or config.get("lora_alpha") != 64:
        raise ValueError("Unexpected rank/alpha: expected r=32, alpha=64")
    adapter_file = adapter / "adapter_model.safetensors"
    if not (adapter / "TRAINING_COMPLETE.json").exists():
        raise ValueError("No completed training marker. Finish training before release export.")
    signature = {"adapter_sha256": sha256(adapter_file), "base": BASE, "revision": BASE_REVISION,
                 "merge": "Original unquantized base plus LoRA delta, fp32 row chunks rounded to fp16"}
    output.mkdir(parents=True, exist_ok=True)
    progress_path = output / "merge_progress.json"
    if progress_path.exists():
        progress = json.loads(progress_path.read_text())
        if progress["signature"] != signature:
            raise ValueError("Existing export belongs to a different adapter or base; choose a new directory")
    else:
        progress = {"signature": signature, "tensors": {}}
    completed = output / "MERGE_COMPLETE.json"
    if completed.exists() and json.loads(completed.read_text()).get("adapter_sha256") == signature["adapter_sha256"]:
        records = progress["tensors"]
        expected_count = json.loads(completed.read_text())["tensor_count"]
        if len(records) == expected_count and all((output / row["file"]).exists() and sha256(output / row["file"]) == row["sha256"] for row in records.values()):
            return output
    base_dir = Path(snapshot_download(BASE, revision=BASE_REVISION, cache_dir=cache,
                        allow_patterns=["*.safetensors", "*.json", "*.txt", "*.model", "*.jinja", "LICENSE*", "NOTICE*"],
                        max_workers=1))
    index = json.loads((base_dir / "model.safetensors.index.json").read_text())
    mapping = index["weight_map"]
    pairs = lora_pairs(adapter_file)
    if set(pairs) - set(mapping):
        raise ValueError(f"Adapter targets missing from original base: {sorted(set(pairs)-set(mapping))}")
    weight_map, total_size = {}, 0
    # Keep one mmap source and one converted tensor alive at a time.
    for i, name in enumerate(sorted(mapping)):
        filename = f"model-{i+1:05d}-of-{len(mapping):05d}.safetensors"
        destination = output / filename
        prior = progress["tensors"].get(name)
        if prior and destination.exists() and sha256(destination) == prior["sha256"]:
            weight_map[name] = filename
            total_size += prior["tensor_bytes"]
            continue
        with safe_open(base_dir / mapping[name], framework="pt", device="cpu") as f:
            weight = f.get_tensor(name)
            merged = merge_tensor(weight, pairs.get(name), config["lora_alpha"] / config["r"])
            tensor_bytes = merged.numel() * merged.element_size()
            temp = destination.with_suffix(".partial")
            save_file({name: merged}, str(temp), metadata={"format": "pt"})
            temp.replace(destination)
            del merged, weight
        gc.collect()
        progress["tensors"][name] = {"file": filename, "tensor_bytes": tensor_bytes, "sha256": sha256(destination)}
        atomic_json(progress_path, progress)
        weight_map[name] = filename
        total_size += tensor_bytes
        print(f"Merged tensor {i+1}/{len(mapping)}", flush=True)
    atomic_json(output / "model.safetensors.index.json", {"metadata": {"total_size": total_size}, "weight_map": weight_map})
    for path in base_dir.iterdir():
        if path.is_file() and path.suffix in {".json", ".txt", ".model", ".jinja"} and path.name != "model.safetensors.index.json" or path.is_file() and path.name.startswith(("LICENSE", "NOTICE")):
            shutil.copyfile(path, output / path.name)
    # Preserve the exact training tokenizer/chat template in all released formats.
    for path in adapter.iterdir():
        if path.is_file() and ("token" in path.name or path.name in {"vocab.json", "merges.txt", "chat_template.jinja"}):
            shutil.copyfile(path, output / path.name)
    model_config = json.loads((output / "config.json").read_text())
    model_config.pop("quantization_config", None)
    model_config["torch_dtype"] = "float16"
    atomic_json(output / "config.json", model_config)
    atomic_json(output / "MERGE_COMPLETE.json", {**signature, "tensor_count": len(weight_map), "total_size": total_size,
                                               "numeric_validation": "Finite tensors and dimension checks; full-model logits still require runtime validation."})
    return output


def run(command, **kwargs):
    subprocess.run([str(x) for x in command], check=True, **kwargs)


def converter_environment(llama, workspace):
    """Overlay converter dependencies without changing the live eval runtime."""
    from importlib.metadata import version
    # The converter pins torch==2.11.0. Reuse the installed CUDA build instead
    # of downloading another multi-gigabyte wheel into this CPU-only export.
    torch_version = version("torch")
    if torch_version.split("+", 1)[0] != "2.11.0":
        raise RuntimeError("Pinned GGUF converter requires installed torch2.11.0; preserve the eval environment and use the provided Colab runtime")
    environment = workspace / "converter-venv"
    # Colab's Python can omit ensurepip. Its existing system pip is visible
    # through system-site-packages, so no bootstrap or Python upgrade is needed.
    # Always rerun creation: a failed ensurepip attempt can leave the config
    # with system-site-packages disabled even though bin/python already exists.
    run([sys.executable, "-m", "venv", "--system-site-packages", "--without-pip", environment])
    python = environment / "bin/python"
    constraints = workspace / "converter-constraints.txt"
    constraints.write_text(f"torch==={torch_version}\n")
    env = dict(os.environ, TMPDIR=str(workspace), PYTHONNOUSERSITE="1",
               PIP_REQUIRE_VIRTUALENV="true", PIP_DISABLE_PIP_VERSION_CHECK="1",
               PIP_CONFIG_FILE=os.devnull,
               OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1",
               NUMEXPR_NUM_THREADS="1")
    env.pop("PYTHONPATH", None)
    env.pop("PYTHONHOME", None)
    for key in ("PIP_TARGET", "PIP_PREFIX", "PIP_ROOT", "PIP_USER"):
        env.pop(key, None)
    isolation_probe = """import importlib.util, sys, sysconfig
from pathlib import Path
expected = Path(sys.argv[1]).resolve()
assert sys.prefix != sys.base_prefix, 'Converter must run in an isolated venv'
assert Path(sys.prefix).resolve() == expected, 'Unexpected converter prefix'
for key in ('purelib', 'platlib', 'scripts'):
    assert Path(sysconfig.get_path(key)).resolve().is_relative_to(expected), 'Package destination escapes converter venv'
assert importlib.util.find_spec('pip') is not None, 'System pip is not visible inside converter venv'
"""
    run([python, "-c", isolation_probe, environment], env=env)
    run([python, "-m", "pip", "install", "--constraint", constraints,
         "-r", llama / "requirements/requirements-convert_hf_to_gguf.txt"], env=env)
    package_probe = """import importlib.metadata as m, json, sys
assert sys.prefix != sys.base_prefix, 'Converter must run in its isolated venv'
names = ['torch', 'numpy', 'transformers', 'sentencepiece', 'protobuf', 'huggingface-hub', 'gguf', 'safetensors', 'tokenizers']
print(json.dumps({'python': sys.version.split()[0], 'packages': {name: m.version(name) for name in names}}))
"""
    versions = json.loads(subprocess.check_output([str(python), "-c", package_probe], env=env, text=True))
    if versions["packages"]["torch"] != torch_version:
        raise RuntimeError("Converter environment unexpectedly replaced the shared torch build")
    return python, env, versions


def gguf(merged, output, workspace):
    output.mkdir(parents=True, exist_ok=True)
    workspace.mkdir(parents=True, exist_ok=True)
    llama = workspace / "llama.cpp"
    if not llama.exists():
        run(["git", "init", llama])
        run(["git", "-C", llama, "remote", "add", "origin", "https://github.com/ggml-org/llama.cpp.git"])
    run(["git", "-C", llama, "fetch", "--depth", "1", "origin", LLAMA_REVISION])
    run(["git", "-C", llama, "checkout", "--detach", LLAMA_REVISION])
    converter_python, temp_env, converter_versions = converter_environment(llama, workspace)
    run(["cmake", "-S", llama, "-B", llama / "build", "-DGGML_CUDA=OFF", "-DLLAMA_CURL=OFF", "-DLLAMA_BUILD_TESTS=OFF", "-DLLAMA_BUILD_SERVER=OFF"])
    run(["cmake", "--build", llama / "build", "--config", "Release", "--target", "llama-quantize", "-j", "1"])
    f16 = workspace / "vibegame-f16.gguf"
    quantized = output / "VibeGame-8B-Q4_K_M.gguf"
    # Converter uses lazy tensor loading and temp-backed writing, so it need not
    # retain the whole fp16 model. No --no-lazy option and no GPU merge allocation.
    run([converter_python, llama / "convert_hf_to_gguf.py", merged, "--outtype", "f16",
         "--outfile", f16, "--use-temp-file", "--model-name", "VibeGame-8B (Qwen2.5-Coder-7B)"], env=temp_env)
    # Upstream defaults to an 8GiB row-buffer budget; reserve host RAM for
    # Colab, the export supervisor, and any completed-run evidence processing.
    run([llama / "build/bin/llama-quantize", "--max-buffer-size", "256",
         f16, quantized, "Q4_K_M", "1"])
    if not quantized.exists() or quantized.stat().st_size < 100_000_000:
        raise RuntimeError("GGUF conversion did not produce a plausible 7B model")
    # Ollama's explicit template mirrors the Qwen chat template used in training.
    (output / "Modelfile").write_text('''FROM ./VibeGame-8B-Q4_K_M.gguf
TEMPLATE """{{- if .Messages }}{{- range .Messages }}<|im_start|>{{ .Role }}
{{ .Content }}<|im_end|>
{{- end }}<|im_start|>assistant
{{ else }}{{ if .System }}<|im_start|>system
{{ .System }}<|im_end|>
{{ end }}{{ if .Prompt }}<|im_start|>user
{{ .Prompt }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ .Response }}{{ end }}"""
PARAMETER stop "<|im_end|>"
PARAMETER stop "<|im_start|>"
PARAMETER num_ctx 4096
PARAMETER num_predict 3072
PARAMETER temperature 0
''')
    atomic_json(output / "GGUF_COMPLETE.json", {"sha256": sha256(quantized), "bytes": quantized.stat().st_size,
                "quantization": "Q4_K_M", "llama_cpp_revision": LLAMA_REVISION,
                "converter_environment": converter_versions,
                "resource_limits": {"merge_threads": 1, "build_jobs": 1, "quantizer_threads": 1,
                                    "quantizer_buffer_mib": 256},
                "source_merge": json.loads((merged / "MERGE_COMPLETE.json").read_text()),
                "runtime_validation": "Not implied by successful conversion; run the frozen evaluation against this artifact."})
    # Only delete our disposable intermediate after the quantized file is complete.
    f16.unlink()
    return quantized


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--adapter", type=Path, required=True)
    p.add_argument("--output", type=Path, default=Path("/content/vibegame-exports"))
    p.add_argument("--cache", type=Path, default=Path("/content/vibegame-base-cache"))
    p.add_argument("--workspace", type=Path, default=Path("/content/vibegame-export-work"))
    p.add_argument("--merge-only", action="store_true")
    args = p.parse_args()
    import psutil
    if psutil.virtual_memory().available < 4 * 1024**3:
        p.error("At least 4GiB available CPU RAM required; terminate the training subprocess before export")
    # Merge and conversion run in stages. A disposable base cache is removed
    # after the verified merge, before GGUF's temporary fp16 files are created.
    for path in (args.output, args.cache, args.workspace):
        path.mkdir(parents=True, exist_ok=True)
    disposable_cache = args.cache / "vibegame-disposable-original-base"
    merge_done = args.output / "merged-fp16/MERGE_COMPLETE.json"
    if not merge_done.exists():
        for path in {os.stat(x).st_dev: x for x in (args.output, args.cache)}.values():
            if shutil.disk_usage(path).free < 34 * 1024**3:
                p.error(f"Need 34GiB free for original-base and fp16 merge staging on {path}; no paid storage is allocated")
    merged = merge(args.adapter, args.output / "merged-fp16", disposable_cache)
    # This nested directory belongs exclusively to this export command.
    if disposable_cache.exists():
        shutil.rmtree(disposable_cache)
    if not args.merge_only:
        for path in {os.stat(x).st_dev: x for x in (args.output, args.workspace)}.values():
            if shutil.disk_usage(path).free < 34 * 1024**3:
                p.error(f"Merged fp16 is safe at {merged}; need 34GiB free on {path} for GGUF temporary files")
        gguf(merged, args.output / "gguf", args.workspace)
    print(f"Exported to {args.output}. Publish/download before the free VM disconnects.")


if __name__ == "__main__":
    main()