"""Vendor the minimal llama.cpp file set this port needs into `vendor/`, so the package is
self-contained (a hub `kernels` build cannot clone at build time).

The file list is the exact transitive dependency set of our three sources, obtained from the
build with:

    cd build && ninja -t deps | grep -oE '/[^ ]*llama\\.cpp/ggml/[^ ]*\\.(cu|cuh|h)'

Re-run this script after bumping the upstream pin; if upstream adds an include, the build fails
loudly with a missing header rather than silently using a stale copy.

Usage: python vendor.py [--src /path/to/llama.cpp] [--rev <git rev>]
"""

import argparse
import os
import shutil
import subprocess

HERE = os.path.dirname(os.path.abspath(__file__))
VENDOR = os.path.join(HERE, "vendor")

FILES = [
    "include/ggml-alloc.h",
    "include/ggml-backend.h",
    "include/ggml-cuda.h",
    "include/ggml.h",
    "include/gguf.h",
    "src/ggml-common.h",
    "src/ggml-impl.h",
    "src/ggml-cuda/common.cuh",
    "src/ggml-cuda/convert.cu",
    "src/ggml-cuda/convert.cuh",
    "src/ggml-cuda/dequantize.cuh",
    "src/ggml-cuda/mma.cuh",
    "src/ggml-cuda/mmq-config-ampere.cuh",
    "src/ggml-cuda/mmq-config-blackwell.cuh",
    "src/ggml-cuda/mmq-config-cdna.cuh",
    "src/ggml-cuda/mmq-config-pascal.cuh",
    "src/ggml-cuda/mmq-config-rdna2.cuh",
    "src/ggml-cuda/mmq-config-rdna3-5.cuh",
    "src/ggml-cuda/mmq-config-rdna3.cuh",
    "src/ggml-cuda/mmq-config-rdna4.cuh",
    "src/ggml-cuda/mmq-load-tiles.cuh",
    "src/ggml-cuda/mmq-vec-dot.cuh",
    "src/ggml-cuda/mmq.cuh",
    "src/ggml-cuda/mmvq.cu",
    "src/ggml-cuda/mmvq.cuh",
    "src/ggml-cuda/mmvf.cu",
    "src/ggml-cuda/mmvf.cuh",
    "src/ggml-cuda/quantize.cu",
    "src/ggml-cuda/quantize.cuh",
    "src/ggml-cuda/unary.cuh",
    "src/ggml-cuda/vecdotq.cuh",
    "src/ggml-cuda/vendors/cuda.h",
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=os.environ.get("LLAMA_CPP_SRC", os.path.join(HERE, "llama.cpp")))
    ap.add_argument("--rev", default=None, help="git rev to check out before copying")
    args = ap.parse_args()

    if not os.path.isdir(args.src):
        subprocess.run(
            ["git", "clone", "https://github.com/ggml-org/llama.cpp.git", args.src], check=True
        )
    if args.rev:
        subprocess.run(["git", "-C", args.src, "checkout", args.rev], check=True)

    rev = subprocess.run(
        ["git", "-C", args.src, "rev-parse", "HEAD"], capture_output=True, text=True, check=True
    ).stdout.strip()

    if os.path.isdir(VENDOR):
        shutil.rmtree(VENDOR)
    for rel in FILES:
        dst = os.path.join(VENDOR, rel)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        shutil.copy2(os.path.join(args.src, "ggml", rel), dst)

    with open(os.path.join(VENDOR, "UPSTREAM"), "w") as f:
        f.write(f"https://github.com/ggml-org/llama.cpp\n{rev}\n")
    shutil.copy2(os.path.join(args.src, "LICENSE"), os.path.join(VENDOR, "LICENSE"))

    print(f"vendored {len(FILES)} files from llama.cpp @ {rev[:12]} into {VENDOR}")


if __name__ == "__main__":
    main()