from __future__ import annotations

import argparse
from pathlib import Path


ORIGINAL_MOE_BLOCK = """            # Prefer to use the MarlinMoE kernel when it is supported.\n            if (\n                not check_moe_marlin_supports_layer(layer, group_size)\n                or current_platform.is_rocm()\n            ):\n"""

GENERIC_MOE_BLOCK = """            # Prefer the generic WNA16 path on CUDA for now. The Marlin MoE\n            # repack op can fail at load time with PTX toolchain mismatches on\n            # some environments even when the layer is otherwise supported.\n            if True:\n"""


def find_vllm_dir() -> Path:
    import vllm  # type: ignore

    return Path(vllm.__file__).resolve().parent


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Patch the active vLLM install for compressed Sarvam inference."
    )
    parser.add_argument(
        "--moe-kernel",
        choices=("generic", "marlin"),
        default="marlin",
        help="Select the WNA16 MoE kernel path to enable in vLLM.",
    )
    return parser.parse_args()


def replace_once(path: Path, old: str, new: str, marker: str) -> None:
    text = path.read_text(encoding="utf-8")
    if marker in text:
        print(f"already patched: {path}")
        return
    if old not in text:
        raise RuntimeError(f"expected block not found in {path}")
    path.write_text(text.replace(old, new, 1), encoding="utf-8")
    print(f"patched {path}")


def replace_either(path: Path, first_old: str, second_old: str, new: str) -> None:
    text = path.read_text(encoding="utf-8")
    if new in text:
        print(f"already patched: {path}")
        return
    if first_old in text:
        path.write_text(text.replace(first_old, new, 1), encoding="utf-8")
        print(f"patched {path}")
        return
    if second_old in text:
        path.write_text(text.replace(second_old, new, 1), encoding="utf-8")
        print(f"patched {path}")
        return
    raise RuntimeError(f"expected block not found in {path}")


def patch_compressed_tensors_moe(vllm_dir: Path, moe_kernel: str) -> None:
    path = (
        vllm_dir
        / "model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py"
    )
    new = GENERIC_MOE_BLOCK if moe_kernel == "generic" else ORIGINAL_MOE_BLOCK
    replace_either(path, ORIGINAL_MOE_BLOCK, GENERIC_MOE_BLOCK, new)


def patch_fused_moe_loader(vllm_dir: Path) -> None:
    path = vllm_dir / "model_executor/layers/fused_moe/fused_moe.py"
    old = """                # If a configuration has been found, return it\n                tuned_config = json.load(f)\n                # Delete triton_version from tuned_config\n                tuned_config.pop(\"triton_version\", None)\n                return {int(key): val for key, val in tuned_config.items()}\n"""
    new = """                # If a configuration has been found, return it\n                tuned_config = json.load(f)\n                if not isinstance(tuned_config, dict):\n                    logger.warning_once(\n                        \"Ignoring malformed MoE tuned config at %s with type %s; \"\n                        \"falling back to the default MoE config.\",\n                        config_file_path,\n                        type(tuned_config).__name__,\n                        scope=\"global\",\n                    )\n                    continue\n                # Delete triton_version from tuned_config\n                tuned_config.pop(\"triton_version\", None)\n                return {int(key): val for key, val in tuned_config.items()}\n"""
    replace_once(
        path, old, new, "Ignoring malformed MoE tuned config at %s with type %s;"
    )


def main() -> None:
    args = parse_args()
    vllm_dir = find_vllm_dir()
    patch_compressed_tensors_moe(vllm_dir, args.moe_kernel)
    patch_fused_moe_loader(vllm_dir)


if __name__ == "__main__":
    main()