import os
from torch_npu.utils._dynamo import (
    _dynamo_register_interface_for_device,
    patch_SkipFunctionVariable,
    patch_TensorVariable_call_method,
)
# All backends need npu/cpu/mps device_op_overrides.
from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system
from .graph import patch_codegen_with_cpp_wrapper
from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream
from ._npu_meta_registration import npu_patch_meta

npu_patch_meta()
_dynamo_register_interface_for_device()
register_device_op_overrides_npu()
patch_has_triton()
patch_is_gpu()
patch_device_supports_tma()
patch_codegen_with_cpp_wrapper()
patch_cache_base_get_system()


def _get_backend() -> str:
    return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")

def _load_mlir_backend():
    import torch
    try:
        import torch_mlir
        from torch_mlir import ir
    except ImportError as err:
        raise ImportError("torch_mlir is not installed, install it first.") from err
    from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin, torch_mlir_patch
    from .lowering_patch import apply_mlir_inductor_patch
    from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import (
        register_mlir_codegen_backend,
    )

    apply_mlir_inductor_patch()
    register_mlir_codegen_backend()


def _load_dvm_backend():
    import torch
    from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin
    from .lowering_patch import apply_mlir_inductor_patch
    from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import (
        register_mlir_codegen_backend,
    )
    apply_mlir_inductor_patch()
    register_mlir_codegen_backend()
    from .dvm import mlir_fusion


def _load_triton_backend():
    import os
    import torch

    has_triton = torch.utils._triton.has_triton()
    if not has_triton:
        import warnings
        warnings.warn("triton-ascend is not installed, install it first.")
        return
    from torch._dynamo.device_interface import (
        get_interface_for_device,
        register_interface_for_device,
    )
    from torch._inductor import lowering as inductor_lowering
    from torch._inductor.choices import InductorChoices
    from torch._inductor.codegen.common import (
        register_backend_for_device,
        register_device_op_overrides,
    )
    from torch._inductor.runtime import autotune_cache
    from torch_npu.npu import device_count
    from torch_npu.utils._dynamo_device import current_device, NpuInterface, set_device
    from torch_npu.utils._inductor import NPUDeviceOpOverrides
    from . import codegen, config as npu_config
    from .codecache import patch_get_cpp_wrapper_header
    from .config import aggresive_autotune, log as npulog, num_vector_core
    from .decomposition import _register_triton_decompositions
    from .lowering import make_reduction, npu_make_fallback
    from .npu_choices import should_use_persistent_reduction
    from .npu_device import NewNPUDeviceOpOverrides
    from .npu_fusion_attention_graph import register_fa_pass
    from .runtime import _load_cached_autotuning
    from .utils import (
        disable_foreach,
        patch_fx_node_is_input_dependent_cudagraph_unsafe,

    )

    from .graph import patch_codegen_with_cpp_wrapper

    def _inductor_register_backend_for_device():
        from .codegen.cpp_wrapper import CppWrapperNpu
        from .codegen.scheduling import NPUTritonScheduling
        from .codegen.wrapper import NPUWrapperCodeGen

        register_backend_for_device(
            "npu", NPUTritonScheduling, NPUWrapperCodeGen, CppWrapperNpu
        )

    _inductor_register_backend_for_device()

    device = get_interface_for_device("npu")

    inductor_lowering.make_reduction = make_reduction
    inductor_lowering.make_fallback = npu_make_fallback

    def patch_torch_for_aoti():
        from .codegen.cpp_utils import patch_device_to_aten
        from .cpp_builder import patch_get_cpp_torch_device_options
        from .fx_passes.joint_graph import patch_constant_fold_uniform_value
        from .utils import patch_is_same_tensor

        patch_get_cpp_torch_device_options()
        patch_is_same_tensor()
        patch_constant_fold_uniform_value()
        patch_device_to_aten()

        from .ir import patch_extern_kernel_codegen_size_asserts

        patch_extern_kernel_codegen_size_asserts()

        patch_get_cpp_wrapper_header()

    if os.environ.get("DISABLE_AOTI_PATCH", "0") != "1":
        patch_torch_for_aoti()

    if npu_config.dump_fx_graph:
        from .codegen.ir_fx import _patch_npu_inductor_ir

        _patch_npu_inductor_ir()

    from .lowering import (
        _enable_full_lowering_fallback,
        _register_npu_inductor_fallbacks,
    )

    _register_triton_decompositions()

    if npu_config.enable_full_lowering_fallback.strip() == "allfallback":
        _enable_full_lowering_fallback()
    else:
        _register_npu_inductor_fallbacks()

    # register fx_pass should be put behind of _register_npu_inductor_decompositions
    def _replace_benchmark_all_configs():
        from torch_npu._compat.inductor import get_CachingAutotuner

        CachingAutotuner = get_CachingAutotuner()
        from .npu_triton_heuristics import benchmark_all_configs

        CachingAutotuner.benchmark_all_configs = benchmark_all_configs

    if aggresive_autotune:
        _replace_benchmark_all_configs()

        os.environ["TRITON_BENCH_METHOD"] = "npu"

    InductorChoices.should_use_persistent_reduction = should_use_persistent_reduction
    autotune_cache._load_cached_autotuning = _load_cached_autotuning

    register_fa_pass()
    disable_foreach()
    patch_fx_node_is_input_dependent_cudagraph_unsafe()
    register_device_op_overrides_npu()


_BACKEND_LOADERS = {
    "mlir": _load_mlir_backend,
    "dvm": _load_dvm_backend,
    "default": _load_triton_backend,
}


def _load_backend():
    from .lowering_patch import restore_inductor_baseline

    # Reset Inductor globals before each backend switch (mlir <-> triton in one process).
    restore_inductor_baseline()

    backend = _get_backend()
    loader = _BACKEND_LOADERS.get(backend, _load_triton_backend)
    loader()

    from .decomposition import _register_shared_decompositions
    _register_shared_decompositions()

    from ..utils._dynamo import _InductorNpuRegistry
    _InductorNpuRegistry._loaded_backend = backend

_load_backend()