#!/usr/bin/env python3
"""Disable ExLlamaV3's x86-only CPU all-reduce sources on pinned ARM64 builds.

Laguna hybrid inference uses vLLM TP2 collectives, not ExLlamaV3's optional
native CPU all-reduce backend. The CUDA EXL3 quantizer, trellis pack/unpack,
reconstruction, and MoE kernels remain compiled. Every upstream x86 source is
hash-gated so a source revision cannot be patched accidentally.
"""

from __future__ import annotations

import hashlib
import platform
import sys
from pathlib import Path


EXPECTED = {
    "avx2_target.cpp": "40e7445f6d25d0ec26185c9c258dc5409291218141e11b22352e96e0f0e9d7d0",
    "avx512_target.cpp": "4f9b51075f769c34884e2b1510328b814b71fa0bb71fb0fdd92b72ebc9ba447a",
    "parallel/all_reduce_cpu_avx2.cpp": "903f14178e4e1451b259fe06d5c3987f86e3c0f5de2de4715998faafd6040b81",
    "parallel/all_reduce_cpu_avx512.cpp": "ccfd0eedfba4ace032bbcc5d2864b8118f30e6447080e4b5cb925f6d082b865b",
}

PACKAGE_INITIALIZERS = {
    "__init__.py": "59430846a04ceba4e33b0ab0233f551b68ec580fb034534c11a965d3cfb42856",
    "modules/__init__.py": "a820ac4a5ab202b8862f7e42a6c21afae78a9653f4e5aaed3b12b271aa69fa21",
    "modules/quant/__init__.py": "cb10d8f118205ce78a35e8c27e322d8cfb2a3bc32be2ca63807437ae055f3df5",
}

MINIMAL_INITIALIZERS = {
    "__init__.py": (
        'from importlib.metadata import version\n\n'
        '__version__ = version("exllamav3")\n'
    ),
    "modules/__init__.py": (
        '"""Minimal module namespace for the Laguna EXL3 encoder image."""\n'
    ),
    "modules/quant/__init__.py": (
        '"""EXL3 quantizer submodules are imported explicitly by the encoder."""\n'
    ),
}

STUB = r"""#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include "avx2_target.h"
#include "avx512_target.h"
#include "parallel/all_reduce_cpu_avx2.h"
#include "parallel/all_reduce_cpu_avx512.h"

bool is_avx2_supported() { return false; }
bool is_avx512_supported() { return false; }
void enable_fast_fp() {}
void enable_fast_fp_avx2() {}
void enable_fast_fp_avx512() {}

[[noreturn]] static void unavailable()
{
    throw std::runtime_error(
        "ExLlamaV3 native x86 CPU all-reduce is unavailable on ARM64; "
        "use the serving runtime's device collective backend");
}

void perform_cpu_reduce(
    PGContext*, size_t, uint32_t, uint8_t*, size_t) { unavailable(); }
void perform_cpu_reduce_avx2(
    PGContext*, size_t, uint32_t, uint8_t*, size_t) { unavailable(); }
void perform_cpu_reduce_avx512(
    PGContext*, size_t, uint32_t, uint8_t*, size_t) { unavailable(); }
void bf16_add_inplace_avx512(
    uint16_t*, const uint16_t*, size_t) { unavailable(); }
"""


def digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> int:
    if len(sys.argv) != 2:
        raise SystemExit("usage: patch_exllamav3_arm64.py EXLLAMAV3_EXT_DIR")
    machine = platform.machine().lower()
    if machine not in {"aarch64", "arm64"}:
        raise RuntimeError(f"refusing ARM64 patch on architecture {machine!r}")
    root = Path(sys.argv[1]).resolve()
    if not root.is_dir():
        raise FileNotFoundError(root)
    for relative, expected in EXPECTED.items():
        source = root / relative
        if digest(source) != expected:
            raise RuntimeError(f"upstream source hash mismatch: {relative}")
        disabled = source.with_suffix(source.suffix + ".x86-disabled")
        if disabled.exists():
            raise FileExistsError(disabled)
        source.rename(disabled)
    package = root.parent
    for relative, expected in PACKAGE_INITIALIZERS.items():
        initializer = package / relative
        if digest(initializer) != expected:
            raise RuntimeError(
                f"upstream package initializer hash mismatch: {relative}"
            )
        initializer.write_text(
            MINIMAL_INITIALIZERS[relative], encoding="utf-8"
        )
    stub = root / "arm64_no_x86.cpp"
    if stub.exists():
        raise FileExistsError(stub)
    stub.write_text(STUB, encoding="utf-8")
    print(
        "ARM64 compatibility applied: disabled four x86-only CPU "
        "all-reduce sources and narrowed optional package imports; "
        "CUDA EXL3 sources retained"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())