"""Map CUDA-style torch code onto the native ``neuron`` device (the counterpart
of ``torch_npu.contrib.transfer_to_npu``).

DeepInfer writes ``device="cuda"``, ``.cuda()``, ``torch.device("cuda")`` and
``torch.cuda.*`` throughout. Instead of touching every call site, this module:

* rewrites ``cuda`` device arguments of torch factory functions / Tensor methods
  to the target device;
* replaces ``torch.device`` with a constructor that maps ``cuda[:N]`` to the
  target device (``isinstance(x, torch.device)`` keeps working);
* replaces ``torch.cuda`` with a proxy backed by ``torch.neuron`` that fills the
  gaps (``mem_get_info``, ``CUDAGraph``, sync-debug mode, ...).

Each process owns exactly one logical NeuronCore, so every ``cuda:N`` maps to
``neuron:0``. With ``DEEPINFER_NEURON_DEVICE=cpu`` everything maps to the host
instead, which runs the pure-torch kernels for correctness work without a device.
"""

from __future__ import annotations

import contextlib
import functools
import types
from collections.abc import Iterator
from typing import Any

import torch
import torch.utils._device  # not imported by `import torch` alone (e.g. TORCH_DEVICE_BACKEND_AUTOLOAD=0)

_REAL_DEVICE = torch.device
_TARGET: torch.device | None = None


def target_device() -> torch.device:
    assert _TARGET is not None, "transfer_to_neuron.install() has not been called"
    return _TARGET


def _map(device: Any) -> Any:
    """Map a cuda device spec to the target; leave everything else untouched."""
    if isinstance(device, str):
        if device == "cuda" or device.startswith("cuda:"):
            return _TARGET
        return device
    if isinstance(device, _REAL_DEVICE):
        return _TARGET if device.type == "cuda" else device
    if isinstance(device, int) and not isinstance(device, bool):
        # Bare ints mean "cuda:<int>" in torch factory calls.
        return _TARGET
    return device


class _DeviceMeta(type):
    def __instancecheck__(cls, instance: Any) -> bool:
        return isinstance(instance, _REAL_DEVICE)

    def __subclasscheck__(cls, subclass: type) -> bool:
        return issubclass(subclass, _REAL_DEVICE)

    # Compare equal to the real class: torch.fx codegen registers the ``device`` global only when the builtin it
    # recorded at import (the real torch.device) ``== torch.device``.
    def __eq__(cls, other: object) -> bool:
        return other is cls or other is _REAL_DEVICE

    def __ne__(cls, other: object) -> bool:
        return not _DeviceMeta.__eq__(cls, other)

    def __hash__(cls) -> int:
        return hash(_REAL_DEVICE)


class _Device(metaclass=_DeviceMeta):
    """Drop-in for ``torch.device`` that maps ``cuda`` to the target device."""

    def __new__(cls, type_or_device: Any = None, index: int | None = None) -> torch.device:  # type: ignore[misc]
        if isinstance(type_or_device, str) and type_or_device.split(":")[0] == "cuda":
            return _TARGET
        if isinstance(type_or_device, _REAL_DEVICE):
            return _map(type_or_device)
        if index is None:
            return _REAL_DEVICE(type_or_device)
        return _REAL_DEVICE(type_or_device, index)


def _wrap_device_kwarg(fn: Any) -> Any:
    @functools.wraps(fn)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        if "device" in kwargs:
            kwargs["device"] = _map(kwargs["device"])
        return fn(*args, **kwargs)

    return wrapper


_FACTORIES = (
    "arange",
    "as_tensor",
    "empty",
    "empty_like",
    "empty_strided",
    "eye",
    "frombuffer",
    "full",
    "full_like",
    "linspace",
    "logspace",
    "ones",
    "ones_like",
    "rand",
    "rand_like",
    "randint",
    "randint_like",
    "randn",
    "randn_like",
    "randperm",
    "range",
    "scalar_tensor",
    "tensor",
    "tril_indices",
    "triu_indices",
    "zeros",
    "zeros_like",
)


def _patch_factories() -> None:
    # `with torch.device(...)` redirects calls whose func is in this lru_cached set; build it from the
    # original C functions before they are wrapped, or the context manager stops applying to them.
    torch.utils._device._device_constructors()
    for name in _FACTORIES:
        if hasattr(torch, name):
            setattr(torch, name, _wrap_device_kwarg(getattr(torch, name)))
    torch.Generator = _wrap_generator(torch.Generator)


def _wrap_generator(gen_cls: type) -> Any:
    class _GeneratorMeta(type):
        def __instancecheck__(cls, instance: Any) -> bool:
            return isinstance(instance, gen_cls)

    class Generator(metaclass=_GeneratorMeta):
        def __new__(cls, device: Any = "cpu") -> Any:  # type: ignore[misc]
            return gen_cls(device=_map(device))

    return Generator


def _patch_tensor_methods() -> None:
    tensor = torch.Tensor
    orig_to = tensor.to

    def to(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
        if args:
            args = (_map(args[0]), *args[1:])
        if "device" in kwargs:
            kwargs["device"] = _map(kwargs["device"])
        return orig_to(self, *args, **kwargs)

    def cuda(self: torch.Tensor, device: Any = None, non_blocking: bool = False, **kwargs: Any) -> torch.Tensor:
        return orig_to(self, _TARGET, non_blocking=non_blocking)

    tensor.to = to
    tensor.cuda = cuda
    for name in ("new_empty", "new_empty_strided", "new_full", "new_ones", "new_tensor", "new_zeros"):
        setattr(tensor, name, _wrap_device_kwarg(getattr(tensor, name)))

    orig_module_to = torch.nn.Module.to

    def module_to(self: torch.nn.Module, *args: Any, **kwargs: Any) -> torch.nn.Module:
        if args:
            args = (_map(args[0]), *args[1:])
        if "device" in kwargs:
            kwargs["device"] = _map(kwargs["device"])
        return orig_module_to(self, *args, **kwargs)

    torch.nn.Module.to = module_to
    torch.nn.Module.cuda = lambda self, device=None: orig_module_to(self, _TARGET)

    orig_set_default_device = torch.set_default_device
    torch.set_default_device = lambda device: orig_set_default_device(_map(device))


class _NullEvent:
    """Host-mode stand-in for a device event: work is already complete when recorded."""

    def __init__(self, enable_timing: bool = False, blocking: bool = False, interprocess: bool = False) -> None:
        self._t: float | None = None

    def record(self, stream: Any = None) -> None:
        import time

        self._t = time.perf_counter()

    def wait(self, stream: Any = None) -> None:
        pass

    def query(self) -> bool:
        return True

    def synchronize(self) -> None:
        pass

    def elapsed_time(self, end: _NullEvent) -> float:
        return ((end._t or 0.0) - (self._t or 0.0)) * 1000.0


class _NullStream:
    def __init__(self, device: Any = None, priority: int = 0, **kwargs: Any) -> None:
        self.cuda_stream = 0

    def wait_stream(self, stream: Any) -> None:
        pass

    def wait_event(self, event: Any) -> None:
        pass

    def record_event(self, event: Any = None) -> Any:
        event = event or _NullEvent()
        event.record()
        return event

    def synchronize(self) -> None:
        pass

    def query(self) -> bool:
        return True

    def __enter__(self) -> _NullStream:
        return self

    def __exit__(self, *exc: Any) -> None:
        pass


class CUDAGraph:
    """Not supported: the Neuron path replaces graph capture with per-bucket ``torch.compile``."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        raise NotImplementedError("CUDA graphs are not available on Neuron; run with DEEPINFER_USE_CUDA_GRAPH=0")


def _patch_cuda_module(host: bool) -> None:
    """Patch ``torch.cuda`` in place (like ``transfer_to_npu``) so internal references stay valid."""
    cuda = torch.cuda
    backend = None if host else torch.neuron

    def mem_get_info(device: Any = None) -> tuple[int, int]:
        if backend is None:
            import psutil

            vm = psutil.virtual_memory()
            return vm.available, vm.total
        total = backend.get_device_properties(0).total_memory
        return total - backend.memory_reserved(), total

    sync_debug_mode = [0]
    funcs: dict[str, Any] = {
        "is_available": lambda: True,
        "is_initialized": lambda: True if backend is None else backend.is_initialized(),
        "device_count": lambda: 1,
        "current_device": lambda: 0,
        "set_device": lambda device=None: None if backend is None else backend.set_device(0),
        "synchronize": lambda device=None: None if backend is None else backend.synchronize(),
        "empty_cache": lambda: None if backend is None else backend.empty_cache(),
        "mem_get_info": mem_get_info,
        # torch.manual_seed itself calls torch.cuda.manual_seed_all, so these must not call back into it;
        # in host mode the CPU generator is the only one and torch.manual_seed already seeds it.
        "manual_seed": lambda seed: None if backend is None else backend.manual_seed(seed),
        "manual_seed_all": lambda seed: None if backend is None else backend.manual_seed_all(seed),
        "is_current_stream_capturing": lambda: False,
        "set_sync_debug_mode": lambda mode: sync_debug_mode.__setitem__(0, mode),
        "get_sync_debug_mode": lambda: sync_debug_mode[0],
        "CUDAGraph": CUDAGraph,
        "graph": CUDAGraph,
        "graph_pool_handle": lambda: None,
        "Event": _NullEvent if backend is None else backend.Event,
        "Stream": _NullStream if backend is None else backend.Stream,
        "current_stream": (lambda device=None: _NullStream()) if backend is None else (lambda device=None: backend.current_stream()),
        "default_stream": (lambda device=None: _NullStream()) if backend is None else (lambda device=None: backend.default_stream()),
        "stream": (lambda s=None: contextlib.nullcontext(s)) if backend is None else backend.stream,
    }
    for name in ("memory_allocated", "max_memory_allocated", "memory_reserved", "max_memory_reserved", "reset_peak_memory_stats"):
        funcs[name] = (lambda *a, **k: 0) if backend is None else getattr(backend, name)
    if backend is None:
        funcs["get_device_properties"] = lambda device=None: types.SimpleNamespace(
            name="neuron-host", total_memory=mem_get_info()[1], multi_processor_count=1, major=0, minor=0
        )
    else:

        def get_device_properties(device: Any = None) -> Any:
            props = backend.get_device_properties(0)
            # CUDA-only fields some callers size buffers with; a logical NeuronCore (LNC=2) has
            # 2 physical cores, reported as the "SM" count.
            return types.SimpleNamespace(
                name=getattr(props, "name", "neuron"),
                total_memory=props.total_memory,
                multi_processor_count=getattr(props, "multi_processor_count", 2),
                major=0,
                minor=0,
            )

        funcs["get_device_properties"] = get_device_properties
    for name, fn in funcs.items():
        setattr(cuda, name, fn)
    cuda.graphs.CUDAGraph = CUDAGraph

    memory = cuda.memory
    memory._record_memory_history = getattr(backend, "_record_memory_history", lambda *a, **k: None)
    memory._dump_snapshot = getattr(backend, "_dump_snapshot", lambda *a, **k: None)


@contextlib.contextmanager
def hide_cuda() -> Iterator[None]:
    """Report ``torch.cuda`` as unavailable, e.g. around torch.compile: dynamo/inductor otherwise take their CUDA
    paths (RNG state save/restore, triton probing, stream capture checks) on the patched module."""
    cuda = torch.cuda
    saved = cuda.is_available
    cuda.is_available = lambda: False
    try:
        yield
    finally:
        cuda.is_available = saved


def install(device: str) -> None:
    """Install the mapping; ``device`` is ``"neuron"`` or ``"cpu"``."""
    global _TARGET
    if _TARGET is not None:
        return
    _TARGET = _REAL_DEVICE(device, 0) if device == "neuron" else _REAL_DEVICE(device)
    _patch_factories()
    _patch_tensor_methods()
    torch.device = _Device
    _patch_cuda_module(host=device != "neuron")