"""MPS Linear Attention - Flash Linear Attention ops for Apple Silicon.

Import this module to patch causal_conv1d and fla for MPS compatibility.
"""

__version__ = "0.1.0"

import sys
import types
import importlib
from functools import lru_cache

from .causal_conv1d import causal_conv1d_fn, causal_conv1d_update
from .delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule


def patch():
    """Patch causal_conv1d and fla modules so transformers detects the fast path."""

    # 1. Create/patch causal_conv1d module
    if "causal_conv1d" not in sys.modules:
        sys.modules["causal_conv1d"] = types.ModuleType("causal_conv1d")
    mod = sys.modules["causal_conv1d"]
    mod.causal_conv1d_fn = causal_conv1d_fn
    mod.causal_conv1d_update = causal_conv1d_update

    # 2. Patch fla.ops.gated_delta_rule
    for name in ("fla", "fla.ops", "fla.ops.gated_delta_rule"):
        if name not in sys.modules:
            sys.modules[name] = types.ModuleType(name)
    gdr = sys.modules["fla.ops.gated_delta_rule"]
    gdr.chunk_gated_delta_rule = chunk_gated_delta_rule
    gdr.fused_recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule

    # 3. Patch fla.modules with a stub FusedRMSNormGated (transformers has its own fallback)
    if "fla.modules" not in sys.modules:
        sys.modules["fla.modules"] = types.ModuleType("fla.modules")
    # FusedRMSNormGated = None means transformers uses its own Qwen3_5RMSNormGated fallback
    sys.modules["fla.modules"].FusedRMSNormGated = None

    # 4. Patch the availability checks in transformers.utils
    try:
        import transformers.utils.import_utils as iu
        # Clear lru_cache and replace functions
        iu.is_causal_conv1d_available = lru_cache(lambda: True)
        iu.is_flash_linear_attention_available = lru_cache(lambda: True)
    except ImportError:
        pass


# Auto-patch on import
patch()
