from __future__ import annotations

import argparse
import importlib
import py_compile
import shutil
import sys
from pathlib import Path

import vllm


BACKUP_SUFFIX = ".qmask_w2a16.bak"


class PatchError(RuntimeError):
    pass


def replace_once(text: str, old: str, new: str, name: str) -> tuple[str, bool]:
    if new in text:
        print(f"  [already patched] {name}")
        return text, False

    if old not in text:
        raise PatchError(f"could not find expected source for patch: {name}")

    print(f"  [patch] {name}")
    return text.replace(old, new, 1), True


def backup_file(path: Path) -> None:
    backup = Path(str(path) + BACKUP_SUFFIX)

    if backup.exists():
        print(f"  [backup exists] {backup}")
        return

    shutil.copy2(path, backup)
    print(f"  [backup] {backup}")


def restore_file(path: Path) -> None:
    backup = Path(str(path) + BACKUP_SUFFIX)

    if not backup.exists():
        print(f"  [no backup] {path}")
        return

    shutil.copy2(backup, path)
    print(f"  [restored] {path}")


def patch_compressed_tensors_moe(path: Path) -> None:
    print(f"\nPatching {path}")

    text = path.read_text()
    original = text

    old = """from vllm.model_executor.layers.quantization.utils.quant_utils import (
    QuantKey,
    kInt4Static32GroupScale,
    kInt4StaticGroupScale,
    kInt8StaticGroupScale,
)
"""

    new = """from vllm.model_executor.layers.quantization.utils.quant_utils import (
    GroupShape,
    QuantKey,
    ScaleDesc,
    kInt4Static32GroupScale,
    kInt4StaticGroupScale,
    kInt8StaticGroupScale,
)
"""

    text, _ = replace_once(text, old, new, "import GroupShape and ScaleDesc")

    old = """        if self.num_bits == 4:
            if self.group_size == 32:
                scale = kInt4Static32GroupScale
            else:
                scale = kInt4StaticGroupScale
        elif self.num_bits == 8:
            assert self.group_size == -1
            scale = kInt8StaticGroupScale
        else:
            raise ValueError(
                "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now."
            )
"""

    new = """        if self.num_bits == 2:
            if not self.symmetric:
                raise ValueError("W2A16 Humming MoE requires symmetric weights.")

            if self.strategy != QuantizationStrategy.GROUP:
                raise ValueError("W2A16 Humming MoE currently requires group quantization.")

            if self.group_size is None or self.group_size <= 0:
                raise ValueError("W2A16 Humming MoE requires a positive group_size.")

            scale = ScaleDesc(torch.float16, True, GroupShape(1, self.group_size))

        elif self.num_bits == 4:
            if self.group_size == 32:
                scale = kInt4Static32GroupScale
            else:
                scale = kInt4StaticGroupScale

        elif self.num_bits == 8:
            assert self.group_size == -1
            scale = kInt8StaticGroupScale

        else:
            raise ValueError(
                f"CompressedTensorsWNA16MoEMethod currently supports int2, int4 and int8; got int{self.num_bits}."
            )
"""

    text, _ = replace_once(text, old, new, "enable W2A16 QuantKey")

    old = "        self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM\n"

    new = (
        "        self.is_transposed = self.wna16_backend not in "
        "(WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.HUMMING)\n"
    )

    text, _ = replace_once(text, old, new, "use N-first layout for Humming")

    old = """        self.moe_kernel = make_wna16_moe_kernel(
            moe_quant_config=self.moe_quant_config,
            moe_config=self.moe,
            experts_cls=self.experts_cls,
            routing_tables=layer._expert_routing_tables(),
            **marlin_args,
        )
"""

    new = """        self.moe_kernel = make_wna16_moe_kernel(
            moe_quant_config=self.moe_quant_config,
            moe_config=self.moe,
            experts_cls=self.experts_cls,
            backend=self.wna16_backend,
            layer=layer,
            routing_tables=layer._expert_routing_tables(),
            **marlin_args,
        )
"""

    text, _ = replace_once(text, old, new, "pass Humming backend and layer to kernel factory")

    old = """    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
        return make_wna16_moe_quant_config(
            w1_scale=layer.w13_weight_scale,
            w2_scale=layer.w2_weight_scale,
            group_size=self.group_size,
            num_bits=self.num_bits,
            w1_zp=getattr(layer, "w13_weight_zero_point", None),
            w2_zp=getattr(layer, "w2_weight_zero_point", None),
            gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
            gemm1_beta=getattr(layer, "swiglu_beta", None),
        )
"""

    new = """    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
        if self.wna16_backend == WNA16MoEBackend.HUMMING:
            from vllm.model_executor.layers.quantization.utils.humming_utils import get_humming_moe_quant_config

            return get_humming_moe_quant_config(
                layer,
                gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
                gemm1_alpha=getattr(layer, "swiglu_alpha", None),
                gemm1_beta=getattr(layer, "swiglu_beta", None),
            )

        return make_wna16_moe_quant_config(
            w1_scale=layer.w13_weight_scale,
            w2_scale=layer.w2_weight_scale,
            group_size=self.group_size,
            num_bits=self.num_bits,
            w1_zp=getattr(layer, "w13_weight_zero_point", None),
            w2_zp=getattr(layer, "w2_weight_zero_point", None),
            gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
            gemm1_beta=getattr(layer, "swiglu_beta", None),
        )
"""

    text, _ = replace_once(text, old, new, "build native Humming MoE quant config")

    if text == original:
        print("  no changes needed")
        return

    backup_file(path)
    path.write_text(text)


def patch_int_wna16(path: Path) -> None:
    print(f"\nPatching {path}")

    text = path.read_text()
    original = text

    old = """    from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
    if isinstance(quant_config, AutoAWQConfig):
"""

    new = """    from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig

    if isinstance(quant_config, QuantizationArgs):
        strategy = getattr(quant_config.strategy, "value", quant_config.strategy)
        qtype = getattr(quant_config.type, "value", quant_config.type)

        return {
            "quant_method": "compressed-tensors",
            "format": "pack-quantized",
            "num_bits": quant_config.num_bits,
            "group_size": quant_config.group_size,
            "strategy": strategy,
            "symmetric": quant_config.symmetric,
            "type": qtype,
        }

    if isinstance(quant_config, AutoAWQConfig):
"""

    text, _ = replace_once(text, old, new, "support compressed-tensors QuantizationArgs in Humming adapter")

    old = """    raise TypeError(
        "Humming WNA16 checkpoint schema requires AutoAWQConfig or "
        "AutoGPTQConfig, "
        f"got {type(quant_config).__name__}."
    )
"""

    new = """    raise TypeError(
        "Humming WNA16 checkpoint schema requires QuantizationArgs, "
        "AutoAWQConfig or AutoGPTQConfig, "
        f"got {type(quant_config).__name__}."
    )
"""

    text, _ = replace_once(text, old, new, "update Humming adapter error message")

    if text == original:
        print("  no changes needed")
        return

    backup_file(path)
    path.write_text(text)


def patch_fused_humming_moe(path: Path) -> None:
    print(f"\nPatching {path}")

    text = path.read_text()
    original = text

    old = """from vllm.platforms import current_platform
from vllm.utils.import_utils import has_humming
"""

    new = """from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from vllm.utils.import_utils import has_humming
"""

    text, _ = replace_once(text, old, new, "import scalar_types")

    old = """    ) -> bool:
        SUPPORTED_W_A = [
"""

    new = """    ) -> bool:
        if weight_key is not None and activation_key is None:
            scale = weight_key.scale

            is_w2a16 = (
                weight_key.dtype == scalar_types.uint2b2
                and weight_key.symmetric
                and scale.static
                and scale.group_shape.row == 1
                and scale.group_shape.col > 0
            )

            if is_w2a16:
                return True

        SUPPORTED_W_A = [
"""

    text, _ = replace_once(text, old, new, "allow symmetric grouped W2A16 Humming MoE")

    if text == original:
        print("  no changes needed")
        return

    backup_file(path)
    path.write_text(text)


def find_vllm_root() -> Path:
    path = Path(vllm.__file__).resolve().parent

    if not path.exists():
        raise PatchError(f"could not find vLLM package directory: {path}")

    return path


def get_target_files(root: Path) -> list[Path]:
    return [
        root / "model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py",
        root / "model_executor/layers/fused_moe/oracle/int_wna16.py",
        root / "model_executor/layers/fused_moe/experts/fused_humming_moe.py",
    ]


def compile_files(files: list[Path]) -> None:
    print("\nSyntax checking patched files...")

    for path in files:
        py_compile.compile(str(path), doraise=True)
        print(f"  [OK] {path}")


def restore(files: list[Path]) -> None:
    print("Restoring backups...")

    for path in files:
        restore_file(path)

    importlib.invalidate_caches()


def parse_args():
    parser = argparse.ArgumentParser(description="Patch vLLM 0.27.x for CT W2A16 Humming MoE.")
    parser.add_argument("--restore", action="store_true", help="Restore .qmask_w2a16.bak files.")
    parser.add_argument("--force", action="store_true", help="Allow patching a vLLM version other than 0.27.x.")
    return parser.parse_args()


def main():
    args = parse_args()

    version = getattr(vllm, "__version__", "unknown")
    root = find_vllm_root()
    files = get_target_files(root)

    print(f"vLLM version: {version}")
    print(f"vLLM directory: {root}")

    for path in files:
        if not path.exists():
            raise PatchError(f"required file does not exist: {path}")

    if args.restore:
        restore(files)
        print("\nRestored.")
        return

    if not str(version).startswith("0.27.") and not args.force:
        raise PatchError(
            f"this patch was written for vLLM 0.27.x, but found {version}. "
            "Use --force only if you checked the source layout."
        )

    patch_compressed_tensors_moe(files[0])
    patch_int_wna16(files[1])
    patch_fused_humming_moe(files[2])

    compile_files(files)
    importlib.invalidate_caches()

    print("\nPatch complete.")
    print("Restart all running vLLM processes before testing.")
    print("For the first test, use --moe-backend humming.")
    print(f"Backups use suffix: {BACKUP_SUFFIX}")


if __name__ == "__main__":
    try:
        main()
    except PatchError as e:
        print(f"\nERROR: {e}", file=sys.stderr)
        sys.exit(1)