from __future__ import annotations

from typing import Dict, List

from .base import DeviceBackend


class CudaBackend(DeviceBackend):
    name = "cuda"
    visibility_env_keys = ("CUDA_VISIBLE_DEVICES",)
    torch_namespace = "cuda"
    flaggems_vendor = "nvidia"
    supports_cuda_graph = True
    supports_mm_cluster_fix = True
    supports_triton_hooks = True
    list_command = (("nvidia-smi", "--list-gpus"), "GPU ")
    # Unchanged from when this list lived in _cudagraph_plugin; CUDA numbers do
    # not move because of the classification refactor.
    fatal_error_markers = (
        "illegal memory access",
        "cudaerrorillegaladdress",
        "out of memory",
        "device-side assert triggered",
        "an illegal instruction",
        "misaligned address",
        "uncorrectable ecc",
    )
    # No CUDA wording is known to poison the context without also matching one
    # of the fatal markers above, so nothing is listed here rather than guessed.
    poisoning_error_markers = ()

    telemetry_argv = (
        "nvidia-smi",
        "--query-gpu=index,clocks.sm,clocks.mem,temperature.gpu,power.draw,utilization.gpu",
        "--format=csv,noheader,nounits",
        "-i", "{dev}",
    )

    def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
        out: List[Dict[str, str]] = []
        for line in raw.splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) != 6 or not parts[0].isdigit():
                continue
            out.append(dict(zip(
                ("index", "sm_clock_mhz", "mem_clock_mhz",
                 "temp_c", "power_w", "util_pct"), parts)))
        return out