"""Registry of zl_bench device backends.

Aliases map common vendor spellings onto the canonical backend name. Unknown
names fail closed rather than defaulting to CUDA: a silent CUDA fallback would
load the cudagraph plugin on the wrong device.
"""

from __future__ import annotations

import os
from typing import Dict, Optional

from .ascend import AscendBackend
from .base import DeviceBackend, parse_device_list
from .cuda import CudaBackend
from .musa import MusaBackend
from .rocm import RocmBackend

ALIASES: Dict[str, str] = {
    "nvidia": "cuda",
    "nv": "cuda",
    "amd": "rocm",
    "hip": "rocm",
    "npu": "ascend",
    "huawei": "ascend",
}

_INSTANCES: Dict[str, DeviceBackend] = {
    "cuda": CudaBackend(),
    "musa": MusaBackend(),
    "rocm": RocmBackend(),
    "ascend": AscendBackend(),
}


class UnknownBackendError(ValueError):
    pass


def known_backends() -> tuple[str, ...]:
    return tuple(sorted(_INSTANCES))


def normalize_backend_name(name: str) -> str:
    key = (name or "").strip().lower()
    if not key:
        return "cuda"
    return ALIASES.get(key, key)


def get_backend(name: Optional[str] = None) -> DeviceBackend:
    raw = name if name is not None else os.environ.get("ZL_BACKEND", "cuda")
    key = normalize_backend_name(raw)
    backend = _INSTANCES.get(key)
    if backend is None:
        raise UnknownBackendError(
            f"unknown backend {raw!r}; known: {', '.join(known_backends())} "
            f"(aliases: {', '.join(sorted(ALIASES))})"
        )
    return backend


def enumerate_devices(backend: Optional[DeviceBackend] = None,
                      explicit: Optional[str] = None) -> list[str]:
    return (backend or get_backend()).enumerate_devices(explicit)


__all__ = [
    "DeviceBackend",
    "UnknownBackendError",
    "enumerate_devices",
    "get_backend",
    "known_backends",
    "normalize_backend_name",
    "parse_device_list",
]