from ._backends.npu import is_available as _backend_is_available
from ._backends.npu import state as npu_state
from ._backends.npu import allocator as npu_allocator
from ._backends.npu.runtime import device_count
from ._backends.npu.streams import Event, Stream
from ._device import device as Device

_MEMORY_FRACTION = None
_NPU_INITIALIZED = False

_default_generators = {}  # device_index -> Generator


def _get_default_generator(device_index=0):
    """Get or create the default NPU generator for a device."""
    if device_index not in _default_generators:
        from ._random import Generator
        gen = Generator(f'npu:{device_index}')
        _default_generators[device_index] = gen
    return _default_generators[device_index]


class _DefaultGeneratorsAccessor:
    """Lazy accessor for default NPU generators, mimicking torch.cuda.default_generators."""
    def __getitem__(self, index):
        return _get_default_generator(index)


default_generators = _DefaultGeneratorsAccessor()



def is_available(verbose=False):
    return _backend_is_available(verbose=verbose)


__all__ = [
    "is_available",
    "device_count",
    "synchronize",
    "current_device",
    "set_device",
    "device",
    "Stream",
    "Event",
    "default_stream",
    "current_stream",
    "stream",
    "set_stream",
    "stream_priority_range",
    "empty_cache",
    "reset_accumulated_memory_stats",
    "reset_peak_memory_stats",
    "memory_stats",
    "max_memory_reserved",
    "memory_reserved",
    "max_memory_allocated",
    "memory_allocated",
    "mem_get_info",
    "memory_summary",
    "memory_snapshot",
    "set_per_process_memory_fraction",
    "get_device_name",
    "get_device_capability",
    "get_device_properties",
    "can_device_access_peer",
    "enable_peer_access",
    "disable_peer_access",
    "pin_memory",
    "is_pinned",
    "is_initialized",
    "init",
    "manual_seed",
    "manual_seed_all",
    "get_rng_state",
    "set_rng_state",
    "get_rng_state_all",
    "set_rng_state_all",
    "default_generators",
    "_get_default_generator",
]


def _normalize_npu_device(device):
    if device is None:
        return Device("npu", index=npu_state.current_device())
    if isinstance(device, Device):
        dev = device
    elif isinstance(device, int):
        dev = Device("npu", index=device)
    else:
        dev = Device(device)
    if dev.type != "npu":
        raise ValueError(f"Expected a cuda device, but got: {dev}")
    if dev.index is None:
        return Device("npu", index=npu_state.current_device())
    return dev




def _get_allocator(device=None):
    dev = _normalize_npu_device(device)
    return npu_allocator.get_allocator(dev.index or 0)


def synchronize(device=None):
    from ._backends.npu import runtime as npu_runtime

    dev = _normalize_npu_device(device)
    runtime = npu_runtime.get_runtime(dev.index or 0)
    _set_initialized()
    if hasattr(runtime, "synchronize_device"):
        runtime.synchronize_device()
    else:
        runtime.synchronize()


def _set_initialized():
    global _NPU_INITIALIZED
    _NPU_INITIALIZED = True


def is_initialized():
    return _NPU_INITIALIZED


def init():
    from ._backends.npu import runtime as npu_runtime

    runtime = npu_runtime.get_runtime(0)
    runtime.init(0)
    _set_initialized()


def _ensure_initialized():
    if not _NPU_INITIALIZED:
        init()


def mem_get_info(device=None):
    _set_initialized()
    dev = _normalize_npu_device(device)
    from ._backends.npu import runtime as npu_runtime

    return npu_runtime.mem_get_info(dev.index or 0)


def stream_priority_range():
    return (0, 0)


def current_device():
    return npu_state.current_device()


def set_device(device):
    _set_initialized()
    dev = _normalize_npu_device(device)
    npu_state.set_device(dev.index or 0)


def default_stream(device=None):
    _set_initialized()
    dev = _normalize_npu_device(device)
    return npu_state.default_stream(dev.index or 0)


def current_stream(device=None):
    _set_initialized()
    dev = _normalize_npu_device(device)
    return npu_state.current_stream(dev.index or 0)


def set_stream(stream):
    npu_state.set_current_stream(stream)


class stream:
    def __init__(self, s):
        self.stream = s
        self._prev = None
        self._dev_ctx = None

    def __enter__(self):
        self._prev = npu_state.current_stream(self.stream.device.index or 0)
        self._dev_ctx = npu_state.device_guard(self.stream.device.index or 0)
        self._dev_ctx.__enter__()
        npu_state.set_current_stream(self.stream)
        return self.stream

    def __exit__(self, exc_type, exc, tb):
        npu_state.set_current_stream(self._prev)
        return self._dev_ctx.__exit__(exc_type, exc, tb)


class device:
    def __init__(self, dev):
        self.dev = _normalize_npu_device(dev)
        self._ctx = None

    def __enter__(self):
        self._ctx = npu_state.device_guard(self.dev.index or 0)
        return self._ctx.__enter__()

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


def memory_allocated(device=None):
    return _get_allocator(device).memory_stats().get("allocated_bytes.all.current", 0)


def max_memory_allocated(device=None):
    return _get_allocator(device).memory_stats().get("allocated_bytes.all.peak", 0)


def memory_reserved(device=None):
    return _get_allocator(device).memory_stats().get("reserved_bytes.all.current", 0)


def max_memory_reserved(device=None):
    return _get_allocator(device).memory_stats().get("reserved_bytes.all.peak", 0)


def memory_stats(device=None):
    return _get_allocator(device).memory_stats()


def reset_peak_memory_stats(device=None):
    return _get_allocator(device).reset_peak_memory_stats()


def reset_accumulated_memory_stats(device=None):
    return _get_allocator(device).reset_accumulated_memory_stats()


def empty_cache(device=None):
    return _get_allocator(device).empty_cache()


def _reset_memory_fraction_for_test():
    global _MEMORY_FRACTION
    _MEMORY_FRACTION = None


def _reset_init_for_test():
    global _NPU_INITIALIZED
    _NPU_INITIALIZED = False


def _get_memory_stats(device=None):
    return {}


def _enforce_memory_fraction(requested_bytes, device=None):
    if _MEMORY_FRACTION is None:
        return
    stats = _get_memory_stats(device)
    total = stats.get("total_reserved_memory", 0)
    if total <= 0:
        return
    reserved = stats.get("total_allocated_memory", 0)
    if reserved + requested_bytes > _MEMORY_FRACTION * total:
        raise RuntimeError("NPU memory fraction exceeded")


def set_per_process_memory_fraction(fraction, device=None):
    if fraction <= 0 or fraction > 1:
        raise ValueError("fraction must be in (0, 1]")
    global _MEMORY_FRACTION
    _MEMORY_FRACTION = float(fraction)


_SOC_NAME = None


def _get_soc_name():
    global _SOC_NAME
    if _SOC_NAME:
        return _SOC_NAME
    from ._backends.npu.acl_loader import ensure_acl
    try:
        acl = ensure_acl()
        soc = acl.get_soc_name()
    except Exception as exc:
        raise RuntimeError(f"Failed to query NPU SoC name: {exc}") from exc
    if isinstance(soc, (bytes, bytearray)):
        soc = soc.decode()
    if not soc:
        raise RuntimeError("Failed to query NPU SoC name")
    _SOC_NAME = str(soc)
    return _SOC_NAME


def _get_device_name(device=None):
    soc = _get_soc_name()
    if soc.lower().startswith("ascend"):
        return soc
    return f"Ascend {soc}"


def _get_device_capability(device=None):
    import re

    soc = _get_soc_name().lower()
    # Capability is a coarse architecture bucket (major/minor), while exact
    # model selection should use get_device_name/get_device_properties().name.
    # 910 generation mapping (project convention):
    # - 910 / 910A / 910B (no numeric suffix): first generation -> (9, 1)
    # - 910B with numeric suffix (for example 910B1/B2/B3/B4): second generation -> (9, 2)
    # - 910C (with or without numeric suffix): third generation -> (9, 3)
    if re.search(r"910c(?:\d+)?", soc):
        return (9, 3)
    if re.search(r"910b\d+", soc):
        return (9, 2)
    if "910a" in soc or "910b" in soc or "910" in soc:
        return (9, 1)
    if "310p" in soc:
        return (3, 1)
    if "310" in soc:
        return (3, 0)
    raise RuntimeError(f"Unknown Ascend SoC: {soc}")




def get_device_name(device=None):
    dev = _normalize_npu_device(device)
    return _get_device_name(dev.index or 0)


def get_device_capability(device=None):
    dev = _normalize_npu_device(device)
    return _get_device_capability(dev.index or 0)


class _DeviceProperties:
    def __init__(self, name, major, minor):
        self.name = name
        self.major = int(major)
        self.minor = int(minor)

    def __repr__(self):
        return (
            f"DeviceProperties(name='{self.name}', "
            f"major={self.major}, minor={self.minor})"
        )


def get_device_properties(device=None):
    name = get_device_name(device)
    major, minor = get_device_capability(device)
    return _DeviceProperties(name=name, major=major, minor=minor)

def can_device_access_peer(device, peer_device):
    return False


def enable_peer_access(peer_device, device=None):
    raise RuntimeError("Peer access is not supported on NPU")


def disable_peer_access(peer_device, device=None):
    raise RuntimeError("Peer access is not supported on NPU")


def pin_memory(tensor):
    return tensor.pin_memory()


def is_pinned(tensor):
    return tensor.is_pinned()


def memory_summary(device=None, abbreviated=False):
    stats = {}
    alloc = _get_allocator(device)
    if alloc is not None and hasattr(alloc, "memory_stats"):
        stats = alloc.memory_stats() or {}
    allocated = stats.get("allocated_bytes.all.current", 0)
    allocated_peak = stats.get("allocated_bytes.all.peak", 0)
    reserved = stats.get("reserved_bytes.all.current", 0)
    reserved_peak = stats.get("reserved_bytes.all.peak", 0)
    active = stats.get("active_bytes.all.current", 0)
    active_peak = stats.get("active_bytes.all.peak", 0)
    inactive = stats.get("inactive_split_bytes.all.current", 0)
    inactive_peak = stats.get("inactive_split_bytes.all.peak", 0)
    lines = [
        "|===========================================================================|",
        "|                             NPU Memory Summary                            |",
        "|===========================================================================|",
        f"| Allocated memory      | current: {allocated:>10} | peak: {allocated_peak:>10} |",
        f"| Reserved memory       | current: {reserved:>10} | peak: {reserved_peak:>10} |",
        f"| Active memory         | current: {active:>10} | peak: {active_peak:>10} |",
        f"| Inactive split memory | current: {inactive:>10} | peak: {inactive_peak:>10} |",
        "|===========================================================================|",
    ]
    if abbreviated:
        return "\n".join(lines[:4])
    return "\n".join(lines)


def memory_snapshot():
    alloc = _get_allocator(None)
    if alloc is not None and hasattr(alloc, "snapshot"):
        return alloc.snapshot()
    return {
        "segments": [],
        "device": current_device() if _backend_is_available() else 0,
        "allocator": "npu",
    }


def manual_seed(seed: int):
    """Set the seed for generating random numbers for the current NPU device."""
    dev_idx = current_device()
    gen = _get_default_generator(dev_idx)
    gen.manual_seed(seed)


def manual_seed_all(seed: int):
    """Set the seed for generating random numbers on all NPU devices."""
    from ._backends.npu.runtime import device_count
    try:
        n = device_count()
    except Exception:
        n = 1
    for i in range(n):
        gen = _get_default_generator(i)
        gen.manual_seed(seed)


def _get_seed(device_index=None):
    """Get current NPU seed for the given device."""
    if device_index is None:
        device_index = current_device()
    return _get_default_generator(device_index)._seed


def _get_and_advance_offset(device_index=None, increment=10):
    """Get (seed, offset) and advance offset for the given device.

    Returns:
        tuple: (seed, offset) as integers.
    """
    if device_index is None:
        device_index = current_device()
    gen = _get_default_generator(device_index)
    return gen.philox_engine_inputs(increment)


def get_rng_state(device=None):
    """Get NPU RNG state as a ByteTensor."""
    dev = _normalize_npu_device(device)
    gen = _get_default_generator(dev.index or 0)
    return gen.get_state()


def set_rng_state(new_state, device=None):
    """Set NPU RNG state from a ByteTensor."""
    dev = _normalize_npu_device(device)
    gen = _get_default_generator(dev.index or 0)
    gen.set_state(new_state)


def get_rng_state_all():
    """Get RNG state for all NPU devices."""
    from ._backends.npu.runtime import device_count
    try:
        n = device_count()
    except Exception:
        n = 1
    return [get_rng_state(i) for i in range(n)]


def set_rng_state_all(states):
    """Set RNG state for all NPU devices."""
    for i, state in enumerate(states):
        set_rng_state(state, device=i)
