import importlib
import importlib.abc
import functools
import inspect
import logging
import os
import sys
import threading
from typing import Any, Optional, TYPE_CHECKING

import torch
import torch_npu
from torch import _TorchCompileWrapper


use_jit_script = False
log = logging.getLogger(__name__)

def _create_npu_autocast_mode_variable(func, args, kwargs):
    from torch._dynamo.variables.base import VariableTracker
    from torch._dynamo.variables.ctx_manager import AutocastModeVariable

    bound_args = inspect.signature(func).bind(*args, **kwargs)
    bound_args.apply_defaults()
    target_values = []
    kwargs.clear()

    for key in ["device_type", "dtype", "enabled", "cache_enabled"]:
        if key == "device_type" and func in [
            torch_npu.npu.amp.autocast,
        ]:
            arg = "npu" if func is torch_npu.npu.amp.autocast else "cpu"
        else:
            arg = bound_args.arguments[key]
        if isinstance(arg, VariableTracker):
            target_values.append(arg.as_python_constant())
        else:
            target_values.append(arg)

    return AutocastModeVariable(target_values, initial_values=None, **kwargs)


def patch_SkipFunctionVariable():
    from torch._dynamo.variables.functions import SkipFunctionVariable
    from torch._dynamo.variables.torch import TorchInGraphFunctionVariable

    def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs):
        if value in [
            torch.npu.stream,
            torch_npu.npu.stream,
            torch_npu.npu.utils.stream,
        ]:
            return TorchInGraphFunctionVariable(value, **kwargs)
        return cls.__new__raw(cls)

    SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__
    SkipFunctionVariable.__new__ = SkipFunctionVariable__new__


def patch_TensorVariable_call_method():
    from torch._dynamo.utils import tensortype_to_dtype
    from torch._dynamo.variables.constant import ConstantVariable
    from torch._dynamo.variables.lists import TupleVariable
    from torch._dynamo.variables.tensor import TensorVariable

    def TensorVariable_call_method(self, tx, name, args, kwargs):
        if (
            name == "type"
            and self.dtype is not None
            and len(args) == 0
            and isinstance(self.device, torch.device)
            and self.device.type == "npu"
        ):
            tensortype = next(
                k for k, v in tensortype_to_dtype.items() if self.dtype in v
            )
            constant_result = ConstantVariable.create(
                f"torch.npu.{tensortype.__name__}"
            )

            if len(args) == 1:
                return constant_result.getitem_const(args[0])
            if args:
                return TupleVariable(
                    [constant_result.getitem_const(a) for a in args]
                )
            return constant_result
        return TensorVariable.call_method_raw(self, tx, name, args, kwargs)

    TensorVariable.call_method_raw = TensorVariable.call_method
    TensorVariable.call_method = TensorVariable_call_method


class _InductorNpuRegistry:
    _disabled_register = False
    _loaded_backend = None

    @classmethod
    def register_inductor_npu(cls):
        if cls._disabled_register:
            return

        current = os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
        if cls._loaded_backend != current:
            if "torch_npu._inductor" not in sys.modules:
                importlib.import_module("torch_npu._inductor")
            else:
                sys.modules["torch_npu._inductor"]._load_backend()
            cls._loaded_backend = current

    @classmethod
    def disable_register(cls):
        cls._disabled_register = True

    @classmethod
    def enable_register(cls):
        cls._disabled_register = False

    @classmethod
    def has_initialized(cls):
        return cls._loaded_backend is not None


def is_inductor_npu_initialized():
    return _InductorNpuRegistry.has_initialized()


def disable_register_inductor_npu():
    _InductorNpuRegistry.disable_register()


def enable_register_inductor_npu():
    _InductorNpuRegistry.enable_register()


def register_inductor_npu():
    _InductorNpuRegistry.register_inductor_npu()


def _resolve_npu_backend(selected_backend=None) -> str:
    """Resolve NPU backend with priority: compile options > config > env."""
    if selected_backend not in (None, "", "default"):
        return selected_backend

    inductor_config = sys.modules.get("torch._inductor.config")
    global_backend = getattr(inductor_config, "npu_backend", None)
    if global_backend not in (None, "", "default"):
        return global_backend

    return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")


def _resolve_npu_backend_from_wrapper(wrapper) -> str:
    return _resolve_npu_backend(wrapper.config.get("npu_backend"))


class _NpuBackendScope:
    """Apply resolved npu backend for one compile invocation and restore env."""

    def __init__(self, backend: str):
        self.backend = backend
        self._old_env = None

    def __enter__(self):
        self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
        try:
            os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
            register_inductor_npu()
            _lazy_inductor_setup()
            if self.backend == "ascendc":
                from torch_npu._inductor.deterministic_cache import (
                    patch_npu_deterministic_level_cache_keys,
                )

                patch_npu_deterministic_level_cache_keys()
        except BaseException:
            self._restore_backend_env()
            raise
        return self

    def __exit__(self, exc_type, exc, tb):
        self._restore_backend_env()
        return False

    def _restore_backend_env(self):
        if self._old_env is None:
            os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
        else:
            os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env


def patch_inductor_wrapper():
    from typing import Any, Optional

    from torch import _TorchCompileInductorWrapper
    from torch.utils._config_module import _ConfigEntry, Config, ConfigModule

    src_init = _TorchCompileInductorWrapper.__init__
    src_get_config_copy = ConfigModule.get_config_copy
    src_call = _TorchCompileInductorWrapper.__call__

    def new_call(self, model_, inputs_):
        backend = _resolve_npu_backend_from_wrapper(self)
        with _NpuBackendScope(backend):
            if backend == "ascendc":
                from torch_npu.dynamo._deterministic_guard import (
                    install_npu_deterministic_level_guard,
                )

                install_npu_deterministic_level_guard()
            return src_call(self, model_, inputs_)

    def new_get_config_copy(self) -> dict[str, Any]:
        ori_dict = src_get_config_copy(self)
        inductor_config = sys.modules.get("torch._inductor.config")
        if inductor_config is None or self is not inductor_config:
            return ori_dict
        if "npu_backend" not in ori_dict:
            ori_dict["npu_backend"] = "default"
            cfg = Config(default="default", value_type=str)
            if "name" in inspect.signature(_ConfigEntry.__init__).parameters:
                self._config["npu_backend"] = _ConfigEntry(cfg, "npu_backend")
            else:
                self._config["npu_backend"] = _ConfigEntry(cfg)

        return ori_dict

    def new_init(self, mode, options, dynamic, name=None):
        self._npu_defer_shape_handling = True
        self._npu_shape_handling_requested = False
        try:
            if name is not None:
                src_init(self, mode, options, dynamic, name)
            else:
                src_init(self, mode, options, dynamic)
        finally:
            del self._npu_defer_shape_handling
            del self._npu_shape_handling_requested
        _lazy_dynamo_setup()
        backend = _resolve_npu_backend_from_wrapper(self)
        if backend == "mlir":
            with _NpuBackendScope(backend):
                log.info("Running MLIR backend")
                device_id = torch_npu.npu.current_device()
                torch_npu._C._recovery_all_npu_stream(device_id)
        if backend == "dvm":
            with _NpuBackendScope(backend):
                log.info("Running dvm backend")

    _TorchCompileInductorWrapper.__call__ = new_call
    _TorchCompileInductorWrapper.__init__ = new_init
    ConfigModule.get_config_copy = new_get_config_copy


def patch_dynamo_optimize():
    from torch_npu.dynamo import _get_global_npu_backend

    src_optimize = torch._dynamo.optimize

    def npu_optimize(*args, **kwargs):
        backend = None
        if "backend" in kwargs:
            backend = kwargs["backend"]
        elif len(args) == 1:
            backend = args[0]

        backend_name = None
        if isinstance(backend, str):
            backend_name = backend
        elif isinstance(backend, _TorchCompileWrapper):
            backend_name = backend.compiler_name

        if backend_name == "npu":
            # Init torchair ahead of running model.
            _get_global_npu_backend(backend_name)
        return src_optimize(*args, **kwargs)

    torch._dynamo.optimize = npu_optimize


def patch_builtin_variable():
    origin_call_id = torch._dynamo.variables.builtin.BuiltinVariable.call_id

    def _wrap_call_id(self, tx, *args):
        if torch._dynamo.variables.builtin.istype(
            args[0], torch._dynamo.variables.streams.EventVariable
        ):
            return torch._dynamo.variables.ConstantVariable.create(id(args[0].value))
        return origin_call_id(self, tx, *args)

    torch._dynamo.variables.builtin.BuiltinVariable.call_id = _wrap_call_id


def patch_stream_event_variable_python_type():
    """
    Preserve backend-specific stream/event Python types in Dynamo.

    PyTorch's generic StreamVariable/EventVariable report torch.Stream and
    torch.Event. NPU subclasses have Python methods that use super(NpuType,
    self), so Dynamo must use the real runtime subclass when tracing those
    methods, especially when profiler wrappers cause Dynamo to inline them.
    """

    def python_type(self):
        return type(self.value)

    streams = torch._dynamo.variables.streams
    streams.StreamVariable.python_type = python_type
    streams.EventVariable.python_type = python_type


def patch_npu_stream_context():
    from torch._dynamo.device_interface import get_interface_for_device
    from torch._dynamo.variables.base import VariableTracker
    from torch._dynamo.variables.streams import StreamContextVariable, StreamVariable
    from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
    if TYPE_CHECKING:
        from torch._dynamo.symbolic_convert import InstructionTranslator

    class NpuStreamContextVariable(StreamContextVariable):
        """This represents NPU stream context with FX graph set_stream node creation."""

        @staticmethod
        def create(
            tx: "InstructionTranslator",
            stream_to_enter: "StreamVariable",
            **kwargs: dict[str, Any],
        ) -> "NpuStreamContextVariable":
            from torch._dynamo.device_interface import get_interface_for_device
            from torch._dynamo.variables.builder import wrap_fx_proxy_cls

            device_interface = get_interface_for_device(stream_to_enter.device)
            current_stream_var = wrap_fx_proxy_cls(
                StreamVariable,
                tx,
                tx.output.create_proxy(
                    "call_function",
                    device_interface.current_stream,
                    (None,),
                    {},
                ),
            )

            return NpuStreamContextVariable(
                stream_to_enter,
                current_stream=current_stream_var,
                device_interface=device_interface,
                **kwargs,
            )

        def __init__(
            self,
            stream: Optional["StreamVariable"],
            current_stream: Optional["StreamVariable"] = None,
            device_interface: Any | None = None,
            **kwargs: Any,
        ) -> None:
            self.current_stream = current_stream
            self.device_interface = device_interface
            super().__init__(stream, **kwargs)

        def enter(
            self, tx: "InstructionTranslator", *args: VariableTracker
        ) -> VariableTracker:
            # Create set_stream node to switch to self.stream
            if self.get_stream():
                tx.output.create_proxy(
                    "call_function",
                    self.device_interface.set_stream,
                    (self.get_stream().as_proxy(),),
                    {},
                )
            return super().enter(tx)

        def exit(
            self, tx: "InstructionTranslator", *args: VariableTracker
        ) -> VariableTracker:
            # First exit the symbolic stream state
            # Create set_stream node to restore current_stream
            if self.get_stream():
                tx.output.create_proxy(
                    "call_function",
                    self.device_interface.set_stream,
                    (self.current_stream.as_proxy(),),
                    {},
                )
            return super().exit(tx, *args)

    def _handle_npu_device_interface_stream(self, tx, stream):
        return NpuStreamContextVariable.create(tx, stream)

    TorchInGraphFunctionVariable._get_handlers()[
        get_interface_for_device("npu").stream
    ] = _handle_npu_device_interface_stream


def patch_npu_current_stream():
    """Reuse PT handle_current_stream so current_stream gets user_object_index."""
    from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
    handlers = TorchInGraphFunctionVariable._get_handlers()
    handlers[torch.npu.current_stream] = handlers[torch.accelerator.current_stream]


def patch_user_defined_class_variable():
    import functools
    from torch._dynamo.variables.user_defined import UserDefinedClassVariable
    from torch._dynamo.variables.torch import TorchCtxManagerClassVariable
    from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
    original_method = UserDefinedClassVariable._in_graph_classes
    class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable):
        def call_function(self, tx, args, kwargs):
            return _create_npu_autocast_mode_variable(self.value, args, kwargs)

    @staticmethod
    @functools.lru_cache(None)
    def patched_in_graph_classes():
        result = original_method()
        result.add(torch.npu.Event)
        result.add(torch.npu.Stream)
        return result

    def UserDefinedClassVariable__new__(cls, value, **kwargs):
        if value in [
            torch.npu.amp.autocast,
            torch_npu.npu.amp.autocast,
            torch.npu.amp.autocast_mode.autocast,
            torch_npu.npu.amp.autocast_mode.autocast,
        ]:
            return NPUTorchCtxManagerClassVariable(value, **kwargs)
        elif value in [
            torch_npu.npu.BoolTensor,
            torch_npu.npu.ByteTensor,
            torch_npu.npu.CharTensor,
            torch_npu.npu.DoubleTensor,
            torch_npu.npu.FloatTensor,
            torch_npu.npu.HalfTensor,
            torch_npu.npu.IntTensor,
            torch_npu.npu.LongTensor,
            torch_npu.npu.ShortTensor,
            torch_npu.npu.BFloat16Tensor,
        ]:
            return TorchInGraphFunctionVariable(value, **kwargs)
        return cls.__new__raw(cls)

    UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes
    UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__
    UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__


def run_once(f):
    """Run a function successfully only once, waiting for concurrent callers."""
    condition = threading.Condition()

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        thread_id = threading.get_ident()
        with condition:
            while wrapper._is_running:
                if wrapper._running_thread == thread_id:
                    return None
                condition.wait()
            if wrapper.has_run:
                return None
            wrapper._is_running = True
            wrapper._running_thread = thread_id

        try:
            result = f(*args, **kwargs)
        except BaseException:
            with condition:
                wrapper._is_running = False
                wrapper._running_thread = None
                condition.notify_all()
            raise

        with condition:
            wrapper.has_run = True
            wrapper._is_running = False
            wrapper._running_thread = None
            condition.notify_all()
        return result

    wrapper.has_run = False
    wrapper._is_running = False
    wrapper._running_thread = None

    def reset_after_fork():
        # The parent thread running f may not exist in the child process.
        nonlocal condition
        condition = threading.Condition()
        wrapper._is_running = False
        wrapper._running_thread = None

    try:
        os.register_at_fork(after_in_child=reset_after_fork)
    except AttributeError:
        pass
    return wrapper


_COMPLETED_DYNAMO_SETUP_STEPS = set()


def _run_dynamo_setup_step(name, setup):
    """Keep successful setup steps idempotent when a later step fails."""
    if name in _COMPLETED_DYNAMO_SETUP_STEPS:
        return
    setup()
    _COMPLETED_DYNAMO_SETUP_STEPS.add(name)


def _find_spec_without_finder(finder, fullname):
    """Delegate to the remaining meta-path finders without bypassing them."""
    try:
        index = sys.meta_path.index(finder)
    except ValueError:
        return importlib.util.find_spec(fullname)

    sys.meta_path.pop(index)
    try:
        return importlib.util.find_spec(fullname)
    finally:
        sys.meta_path.insert(min(index, len(sys.meta_path)), finder)


class _DynamoPostImportLoader(importlib.abc.Loader):
    def __init__(self, loader, finder):
        self._loader = loader
        self._finder = finder

    def create_module(self, spec):
        create_module = getattr(self._loader, "create_module", None)
        return create_module(spec) if create_module is not None else None

    def exec_module(self, module):
        self._loader.exec_module(module)
        _lazy_dynamo_setup()
        if self._finder in sys.meta_path:
            sys.meta_path.remove(self._finder)


class _DynamoPostImportFinder(importlib.abc.MetaPathFinder):
    _target = "torch._dynamo"

    def find_spec(self, fullname, path=None, target=None):
        if fullname != self._target:
            return None
        spec = _find_spec_without_finder(self, fullname)
        if spec is not None and spec.loader is not None:
            spec.loader = _DynamoPostImportLoader(spec.loader, self)
        return spec


def _install_dynamo_post_import_trigger():
    """Set up NPU integration whenever Dynamo is first imported."""
    if "torch._dynamo" in sys.modules:
        _lazy_dynamo_setup()
        return
    if not any(isinstance(finder, _DynamoPostImportFinder) for finder in sys.meta_path):
        sys.meta_path.insert(0, _DynamoPostImportFinder())


@run_once
def add_dynamo_methods_init():
    steps = (
        ("device_interface", _dynamo_register_interface_for_device),
        ("skip_function_variable", patch_SkipFunctionVariable),
        ("tensor_variable", patch_TensorVariable_call_method),
        ("user_defined_class_variable", patch_user_defined_class_variable),
        ("stream_event_variable", patch_stream_event_variable_python_type),
        ("npu_stream_context", patch_npu_stream_context),
        ("npu_current_stream", patch_npu_current_stream),
        ("builtin_variable", patch_builtin_variable),
    )
    for name, setup in steps:
        _run_dynamo_setup_step(name, setup)


@functools.lru_cache(None)
def has_triton() -> bool:
    from torch.utils._triton import has_triton_package

    if not has_triton_package():
        return False

    from torch._dynamo.device_interface import get_interface_for_device

    def cuda_extra_check(device_interface):
        return True

    def cpu_extra_check(device_interface):
        import triton.backends

        return "cpu" in triton.backends.backends

    def _return_true(device_interface):
        return True

    triton_supported_devices = {
        "cuda": cuda_extra_check,
        "xpu": _return_true,
        "cpu": cpu_extra_check,
        "npu": _return_true,
    }

    def is_device_compatible_with_triton():
        _dynamo_register_interface_for_device()
        for device, extra_check in triton_supported_devices.items():
            device_interface = get_interface_for_device(device)
            if device_interface.is_available() and extra_check(device_interface):
                return True
        return False

    return is_device_compatible_with_triton()


def patch_has_triton():
    from torch.utils import _triton

    _triton.has_triton = has_triton


@run_once
def _inject_inductor_npu_backend_config():
    """Inject NPU entries into torch._inductor.config on first use."""
    torch._inductor.config.get_config_copy()


@run_once
def _lazy_dynamo_setup():
    """Initialize the Dynamo integration on the first graph-capture operation."""
    add_dynamo_methods_init()

    from torch_npu.dynamo import _register_backends
    _run_dynamo_setup_step("backends", _register_backends)

    from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
    _run_dynamo_setup_step("trace_rules", _patch_npu_trace_rules)

    _run_dynamo_setup_step("dynamo_optimize", patch_dynamo_optimize)


@run_once
def _lazy_inductor_setup():
    """Initialize NPU Inductor support only for an Inductor-based backend."""
    register_inductor_npu()

    from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
    _apply_npugraph_tree_methods()

    _inject_inductor_npu_backend_config()


@run_once
def install_npugraph_mark_step_trigger():
    """Expose the public NPUGraph step API without importing compiler internals."""
    def npugraph_mark_step_begin():
        from torch_npu.npu._graph_tree_state import mark_step_begin
        return mark_step_begin()

    torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin


@run_once
def _dynamo_register_interface_for_device():
    from torch._dynamo.device_interface import register_interface_for_device
    from torch_npu.utils._dynamo_device import NpuInterface

    register_interface_for_device("npu", NpuInterface)
    for i in range(32):
        register_interface_for_device(f"npu:{i}", NpuInterface)


def add_dynamo_methods():
    patch_has_triton()

    from torch_npu.dynamo import _install_lazy_torchair

    _install_lazy_torchair()
    _install_dynamo_post_import_trigger()
    if "npugraph_ex" not in sys.modules:
        from torch_npu.dynamo import _LazyNpuGraphEx
        sys.modules["npugraph_ex"] = _LazyNpuGraphEx("npugraph_ex")
    patch_inductor_wrapper()
    install_npugraph_mark_step_trigger()