"""AWS Neuron (Trainium) discovery using device files and ``neuron-ls``.

One DeepInfer device is one logical NeuronCore; visibility is controlled by
``NEURON_RT_VISIBLE_CORES``. A trn2.48xlarge exposes 16 devices x 4 logical
cores (LNC=2) = 64 devices.
"""

from __future__ import annotations

import functools
import json
import re
import subprocess
from pathlib import Path

_DEVICE_PATTERN = re.compile(r"^neuron(\d+)$")


def _path(root: Path, absolute: str) -> Path:
    return root / absolute.lstrip("/")


def _device_node_indices(*, _root: Path = Path("/")) -> tuple[int, ...]:
    device_dir = _path(_root, "/dev")
    if not device_dir.is_dir():
        return ()
    try:
        return tuple(sorted(int(m.group(1)) for p in device_dir.iterdir() if (m := _DEVICE_PATTERN.match(p.name)) is not None))
    except OSError as error:
        raise OSError(f"cannot enumerate {device_dir}: {error}") from error


def is_available(*, _root: Path = Path("/")) -> bool:
    return bool(_device_node_indices(_root=_root))


@functools.lru_cache(maxsize=1)
def _neuron_ls() -> tuple[dict, ...]:
    try:
        output = subprocess.check_output(["neuron-ls", "-j"], text=True, stderr=subprocess.DEVNULL)
    except (FileNotFoundError, subprocess.CalledProcessError):
        return ()
    return tuple(json.loads(output))


def discover_cores(*, _root: Path = Path("/")) -> list[tuple[int, int, int]]:
    """Logical NeuronCores as ``(core_id, device_id, numa_node)``; NUMA is -1 when unknown."""
    cores: list[tuple[int, int, int]] = []
    if _root == Path("/"):
        for device in _neuron_ls():
            numa = int(device.get("numa_node", -1))
            cores.extend((int(core), int(device["neuron_device"]), numa) for core in device["neuroncore_ids"])
    if cores:
        return sorted(cores)
    # Fallback without neuron-ls: sysfs core_count is physical cores; trn2 defaults to LNC=2.
    lnc = 2
    for index in _device_node_indices(_root=_root):
        count_path = _path(_root, f"/sys/class/neuron_device/neuron{index}/core_count")
        physical = int(count_path.read_text().strip()) if count_path.is_file() else 8
        base = len(cores)
        cores.extend((base + i, index, -1) for i in range(physical // lnc))
    return cores


def physical_device_count(*, _root: Path = Path("/")) -> int:
    return len(discover_cores(_root=_root))


def device_numa_node(index: int, *, _root: Path = Path("/")) -> int:
    cores = discover_cores(_root=_root)
    if index >= len(cores):
        raise IndexError(f"no neuron core at index {index}")
    numa = cores[index][2]
    if numa < 0:
        raise RuntimeError(f"NUMA metadata is unavailable for neuron core {index}")
    return numa


__all__ = ["device_numa_node", "discover_cores", "is_available", "physical_device_count"]