"""Monkey-patch entry point for applying Liger kernels to Megatron-Core."""

from __future__ import annotations

import importlib
import logging

logger = logging.getLogger(__name__)

_PATCH_MARKER = "__liger_patched__"

# Megatron modules that bind a cross-entropy function into their own namespace with
# ``from ... import``. Replacing the function on its defining module does not reach a copy
# that was bound before the patch ran, so these are rebound as well. Each entry exists
# only on some Megatron versions; missing modules or attributes are skipped.
_FUSED_CE_CALL_SITES = (
    # Released megatron-core (0.19 and earlier): LanguageModule.compute_language_model_loss
    # calls this module global.
    ("megatron.core.models.common.language_module.language_module", "fused_vocab_parallel_cross_entropy"),
    # Megatron-LM main (NVIDIA/Megatron-LM#6766): select_cross_entropy returns this module
    # global, and LanguageModule caches the result when the model is built.
    ("megatron.core.models.backends", "_fused_ce"),
)
_UNFUSED_CE_CALL_SITES = (
    # Released megatron-core: compute_language_model_loss calls the package-level re-export
    # ``tensor_parallel.vocab_parallel_cross_entropy``, which is bound when the package is
    # imported, i.e. always before its ``cross_entropy`` submodule can be patched.
    ("megatron.core.tensor_parallel", "vocab_parallel_cross_entropy"),
)


def apply_liger_kernel_to_megatron(
    rms_norm: bool = True,
    cross_entropy: bool = False,
    swiglu: bool = False,
) -> None:
    """Patch Megatron-Core to use Liger Triton kernels.

    Idempotent. Targets Megatron's ``BackendSpecProvider``,
    ``transformer_block.LayerNormImpl``, and (optionally) both of Megatron's
    vocab-parallel cross-entropy entry points so models that route through
    the standard spec system pick up Liger without per-model code.

    Args:
        rms_norm: When ``True`` (default) replace both
            ``LocalSpecProvider.layer_norm`` (per-layer norm slots) and
            ``transformer_block.LayerNormImpl`` (the block-level
            ``final_layernorm`` slot) so all RMSNorm modules in the model
            become ``LigerMegatronRMSNorm``.
        cross_entropy: When ``True`` replace both
            ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``
            (fused path) and
            ``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``
            (unfused path) with Liger's Triton cross-entropy, along with the
            copies Megatron's call sites bound at import time. Default
            ``False`` so adopters opt in explicitly. All tensor-parallel sizes
            are supported (TP=1 and TP>1 share the same vocab-parallel
            kernel). The fused wrapper matches native's
            ``(logits, target, tp_group)`` signature exactly; the unfused
            wrapper additionally honors a runtime ``label_smoothing``
            argument, matching native's
            ``(logits, target, label_smoothing=0.0, tp_group=None)``.
        swiglu: When ``True`` replace
            ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger's
            Triton SiLU-multiply kernel, covering the dense ``MLP`` and the MoE
            ``SharedExpertMLP``. Default ``False`` so adopters opt in explicitly.
            See ``_patch_swiglu_function`` for the configurations this reaches
            and the ones that stay on Megatron.

    Notes:
        Call this BEFORE building your model. Patching after instantiation
        will not retroactively swap modules already created.

        The RMSNorm patches only affect the local (non-TE) backend. Mixing
        Liger norms with ``TESpecProvider`` requires a custom
        ``BackendSpecProvider`` subclass because TE's
        ``TELayerNormColumnParallelLinear`` folds the norm into the QKV
        linear; naive substitution would either double-norm or skip the norm.

        For explicit kernel configuration (custom ``ignore_index``,
        ``label_smoothing``, etc.) instantiate ``LigerMegatronCrossEntropy``
        directly and wire it into your model (Mode 2). The monkey-patch path
        is intentionally a transparent drop-in: it matches Megatron's native
        defaults so callers can flip Liger on without touching loss config.
    """
    if rms_norm:
        _patch_local_spec_provider_layer_norm()
        _patch_transformer_block_layernorm_impl()
    if cross_entropy:
        _patch_fused_vocab_parallel_cross_entropy()
        _patch_vocab_parallel_cross_entropy()
    if swiglu:
        _patch_swiglu_function()


def _patch_local_spec_provider_layer_norm() -> None:
    from megatron.core.models import backends

    from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm

    if getattr(backends.LocalSpecProvider.layer_norm, _PATCH_MARKER, False):
        return  # already patched

    original_layer_norm = backends.LocalSpecProvider.layer_norm

    def patched_layer_norm(self, rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False, **kwargs):
        if rms_norm:
            return LigerMegatronRMSNorm
        return original_layer_norm(self, rms_norm=rms_norm, for_qk=for_qk, has_residual=has_residual, **kwargs)

    setattr(patched_layer_norm, _PATCH_MARKER, True)
    setattr(patched_layer_norm, "__wrapped__", original_layer_norm)
    backends.LocalSpecProvider.layer_norm = patched_layer_norm

    logger.info(
        "Patched megatron.core.models.backends.LocalSpecProvider.layer_norm "
        "to return LigerMegatronRMSNorm for rms_norm=True."
    )


def _patch_transformer_block_layernorm_impl() -> None:
    """Cover the block-level ``final_layernorm`` slot.

    ``TransformerBlock`` uses a module-level global ``LayerNormImpl`` (chosen
    at import time from TE / Apex / WrappedTorchNorm) to fill the
    ``final_layernorm`` slot when the caller passes a per-layer spec rather
    than a ``TransformerBlockSubmodules``. Our spec-provider patch only
    catches the per-layer slots, so without this second patch the trailing
    norm stays as PyTorch's ``nn.RMSNorm``.

    We only displace the ``WrappedTorchNorm`` fallback. Users on TE or Apex
    chose those deliberately; replacing them would undo TE's LN+Linear
    fusion or surprise Apex users.
    """
    from megatron.core.transformer import transformer_block
    from megatron.core.transformer.torch_norm import WrappedTorchNorm

    from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm

    if getattr(transformer_block.LayerNormImpl, _PATCH_MARKER, False):
        return

    original = transformer_block.LayerNormImpl
    if original is not WrappedTorchNorm:
        logger.info(
            "transformer_block.LayerNormImpl is %s; skipping block-level patch "
            "(Liger only displaces the pure-torch fallback).",
            getattr(original, "__name__", str(original)),
        )
        return

    class _LigerOrTorchNorm:
        """Routes to ``LigerMegatronRMSNorm`` when ``config.normalization``
        is ``"RMSNorm"``; otherwise instantiates the original implementation.

        Keeps LayerNorm users on PyTorch's ``nn.LayerNorm`` while routing
        RMSNorm users through Liger for the final norm slot.
        """

        def __new__(cls, config, hidden_size, eps=1e-5, **kwargs):
            if getattr(config, "normalization", None) == "RMSNorm":
                return LigerMegatronRMSNorm(config=config, hidden_size=hidden_size, eps=eps, **kwargs)
            return original(config=config, hidden_size=hidden_size, eps=eps, **kwargs)

    setattr(_LigerOrTorchNorm, _PATCH_MARKER, True)
    setattr(_LigerOrTorchNorm, "__wrapped__", original)
    transformer_block.LayerNormImpl = _LigerOrTorchNorm

    logger.info(
        "Patched megatron.core.transformer.transformer_block.LayerNormImpl "
        "to route RMSNorm configs through LigerMegatronRMSNorm."
    )


# Sentinel for "caller did not pass this kwarg". A plain ``0.0`` default would be
# observationally indistinguishable from "user explicitly asked for 0.0" — and Megatron's
# native vocab_parallel_cross_entropy accepts call-time ``label_smoothing=0.0`` as a real
# request for that value. We must not silently override.
_LABEL_SMOOTHING_UNSET = object()


def _rebind_call_sites(call_sites, original, replacement) -> bool:
    """Point every call site still bound to ``original`` at ``replacement``.

    Importing a call-site module here is harmless: the defining module is already patched,
    so a first import binds ``replacement`` directly. A call site bound to some other
    function (another framework's override) is left alone.

    Returns whether any call site now holds ``replacement``.
    """
    reached = False
    for module_name, attr in call_sites:
        try:
            module = importlib.import_module(module_name)
        except ImportError:
            continue
        if getattr(module, attr, None) is original:
            setattr(module, attr, replacement)
        reached = reached or getattr(module, attr, None) is replacement
    return reached


def _patch_fused_vocab_parallel_cross_entropy() -> None:
    """Replace ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``.

    Wraps a single ``LigerMegatronCrossEntropy`` instance (constructed with class defaults
    that match Megatron's native fused-CE behavior) in a closure matching Megatron's fused-CE
    signature ``(logits, target, tp_group)``. Idempotent: a sentinel attribute on the
    replacement prevents wrappers from stacking.

    Megatron never calls the function through this module attribute: each call site binds
    its own copy with ``from ... import`` (see ``_FUSED_CE_CALL_SITES``). Those copies are
    rebound too, so the patch works regardless of import order. Logs a warning when no
    known call site ends up on Liger.
    """
    try:
        import megatron.core.fusions.fused_cross_entropy as fused_ce
    except ImportError as exc:
        raise ImportError(
            "apply_liger_kernel_to_megatron(cross_entropy=True) requires megatron-core to be "
            "installed. Expected symbol path: "
            "megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy."
        ) from exc

    if not hasattr(fused_ce, "fused_vocab_parallel_cross_entropy"):
        raise ImportError(
            "megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy not "
            "found. The symbol path may have changed in your Megatron-LM version. Please file "
            "an issue on https://github.com/linkedin/Liger-Kernel with your megatron-core version."
        )

    if getattr(fused_ce.fused_vocab_parallel_cross_entropy, _PATCH_MARKER, False):
        return  # already patched

    original = fused_ce.fused_vocab_parallel_cross_entropy

    from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy

    ce = LigerMegatronCrossEntropy()

    def liger_fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target, tp_group=None):
        return ce(vocab_parallel_logits, target, tp_group=tp_group)

    setattr(liger_fused_vocab_parallel_cross_entropy, _PATCH_MARKER, True)
    setattr(liger_fused_vocab_parallel_cross_entropy, "__wrapped__", original)
    fused_ce.fused_vocab_parallel_cross_entropy = liger_fused_vocab_parallel_cross_entropy
    reached = _rebind_call_sites(_FUSED_CE_CALL_SITES, original, liger_fused_vocab_parallel_cross_entropy)

    logger.info(
        "Patched megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy with Liger cross-entropy."
    )
    if not reached:
        logger.warning(
            "Could not find where this Megatron version calls fused_vocab_parallel_cross_entropy; "
            "configs with cross_entropy_loss_fusion=True may keep using Megatron's kernel. Please file "
            "an issue on https://github.com/linkedin/Liger-Kernel with your megatron-core version."
        )


def _patch_vocab_parallel_cross_entropy() -> None:
    """Replace ``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``.

    This is Megatron's *unfused* eager-Python vocab-parallel CE path, dispatched to when
    ``config.cross_entropy_loss_fusion=False``. Its signature accepts ``label_smoothing``
    at call time, so the wrapper honors a runtime value when the caller actually passed
    one. A sentinel disambiguates "caller passed 0.0" (use 0.0) from "caller didn't pass"
    (use class default).

    Released megatron-core calls the package-level re-export instead, which is rebound too
    (see ``_UNFUSED_CE_CALL_SITES``). Megatron-LM main looks the function up on this module
    at call time, so the module attribute alone reaches it.
    """
    try:
        import megatron.core.tensor_parallel.cross_entropy as unfused_ce
    except ImportError as exc:
        raise ImportError(
            "apply_liger_kernel_to_megatron(cross_entropy=True) requires megatron-core to be "
            "installed. Expected symbol path: "
            "megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy."
        ) from exc

    if not hasattr(unfused_ce, "vocab_parallel_cross_entropy"):
        raise ImportError(
            "megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy not "
            "found. The symbol path may have changed in your Megatron-LM version. Please file "
            "an issue on https://github.com/linkedin/Liger-Kernel with your megatron-core version."
        )

    if getattr(unfused_ce.vocab_parallel_cross_entropy, _PATCH_MARKER, False):
        return  # already patched

    original = unfused_ce.vocab_parallel_cross_entropy

    from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy

    # Class-default instance; reused for every call where the caller doesn't pass
    # label_smoothing. Avoids allocating a fresh module per CE call in the common case
    # (Megatron's own LanguageModule.compute_language_model_loss dispatch does not pass
    # label_smoothing — it always lands here).
    default_ce = LigerMegatronCrossEntropy()

    def liger_vocab_parallel_cross_entropy(
        vocab_parallel_logits,
        target,
        label_smoothing=_LABEL_SMOOTHING_UNSET,
        tp_group=None,
    ):
        # Sentinel-based "did the caller pass this?" check so that an explicit
        # label_smoothing=0.0 from the caller is honored verbatim (matching Megatron's
        # native vocab_parallel_cross_entropy contract). Construct a fresh
        # LigerMegatronCrossEntropy only on the runtime-override path; nn.Module
        # construction is microseconds vs. CE-kernel milliseconds.
        if label_smoothing is _LABEL_SMOOTHING_UNSET:
            return default_ce(vocab_parallel_logits, target, tp_group=tp_group)
        ce = LigerMegatronCrossEntropy(label_smoothing=label_smoothing)
        return ce(vocab_parallel_logits, target, tp_group=tp_group)

    setattr(liger_vocab_parallel_cross_entropy, _PATCH_MARKER, True)
    setattr(liger_vocab_parallel_cross_entropy, "__wrapped__", original)
    unfused_ce.vocab_parallel_cross_entropy = liger_vocab_parallel_cross_entropy
    _rebind_call_sites(_UNFUSED_CE_CALL_SITES, original, liger_vocab_parallel_cross_entropy)

    logger.info(
        "Patched megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy with Liger cross-entropy."
    )


def _patch_swiglu_function() -> None:
    """Replace ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger.

    Covers ``MLP.forward`` and ``SharedExpertMLP.forward`` when
    ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True``
    and ``config.activation_func is F.silu``.

    Megatron's ``bias_swiglu_impl`` is intentionally left alone: it is a plain function
    that resolves ``SwiGLUFunction`` from its own globals on every call, so replacing the
    class reaches every caller -- including modules that already did
    ``from ... import bias_swiglu_impl`` -- regardless of import order. Nothing outside
    ``fused_bias_swiglu`` references ``SwiGLUFunction``, so there is exactly one binding to
    replace and no stale copies can exist.

    Patching one layer further down (the ``swiglu`` / ``swiglu_back`` math helpers) is not
    an option: they are ``@jit_fuser``-decorated, and ``jit_fuser`` is ``torch.jit.script``
    below torch 2.2, which compiles them at import time -- the patch would silently no-op.

    Not patched: ``BiasSwiGLUFunction`` (Liger's kernel has no bias term, so a non-``None``
    bias keeps using Megatron by construction), ``WeightedSwiGLUFunction`` (MoE routed
    experts, needs routing-weight grad), and ``config.use_te_activation_func=True``
    (TransformerEngine owns that path).
    """
    try:
        import megatron.core.fusions.fused_bias_swiglu as fused_swiglu
    except ImportError as exc:
        raise ImportError(
            "apply_liger_kernel_to_megatron(swiglu=True) requires megatron-core to be "
            "installed. Expected symbol path: "
            "megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction."
        ) from exc

    if not hasattr(fused_swiglu, "SwiGLUFunction"):
        raise ImportError(
            "megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction not found. The symbol "
            "path may have changed in your Megatron-LM version. Please file an issue on "
            "https://github.com/linkedin/Liger-Kernel with your megatron-core version."
        )

    if getattr(fused_swiglu.SwiGLUFunction, _PATCH_MARKER, False):
        return  # already patched

    original = fused_swiglu.SwiGLUFunction

    from liger_kernel.ops.swiglu import LigerFusedGateUpSiLUMulFunction

    # Deduplicate fallback logs so unsupported configs do not spam every step.
    logged_fallbacks = set()

    class _LigerSwiGLUFunction:
        """Adapter matching Megatron's ``SwiGLUFunction.apply`` signature.

        This is not an ``autograd.Function``. It exists to map Megatron's
        ``(input, fp8_input_store, cpu_offload_input)`` call shape to Liger's kernel
        signature without changing call sites.
        """

        @staticmethod
        def apply(input, fp8_input_store=False, cpu_offload_input=False):
            # FP8 input-store and CPU offload use Megatron-specific backward storage
            # semantics, so these paths defer to native.
            reason = None
            if fp8_input_store:
                reason = "config.activation_func_fp8_input_store=True"
            elif cpu_offload_input:
                reason = "CPU activation offloading enabled"
            if reason is not None:
                if reason not in logged_fallbacks:
                    logged_fallbacks.add(reason)
                    logger.info(
                        "Liger SwiGLU is deferring to Megatron's native SwiGLUFunction: %s. "
                        "Numerics and memory behavior are unchanged for this configuration.",
                        reason,
                    )
                return original.apply(input, fp8_input_store, cpu_offload_input)

            return LigerFusedGateUpSiLUMulFunction.apply(input, False)

    setattr(_LigerSwiGLUFunction, _PATCH_MARKER, True)
    setattr(_LigerSwiGLUFunction, "__wrapped__", original)
    fused_swiglu.SwiGLUFunction = _LigerSwiGLUFunction

    logger.info("Patched megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction with Liger SwiGLU.")
