from __future__ import annotations

from typing import Dict, List

from .base import DeviceBackend


class MusaBackend(DeviceBackend):
    name = "musa"
    visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
    torch_namespace = "musa"
    flaggems_vendor = "musa"
    # torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so
    # capture works once the timing path takes them from Triton's device
    # interface instead of torch.cuda. Measured on an S5000 over 15 operators:
    # 11 captured, 3 fell back cleanly with a marker (dropout/rand on the RNG
    # seed, sort on a capture-illegal stream op), 1 poisoned the context (see
    # capture_deny_ops). Capture is worth having because it removes launch
    # overhead the direct timer cannot: batch_norm 16x16x8x48 went 0.0758ms ->
    # 0.0089ms, tanh/gelu 4096x4096 about 3.3x, while GEMMs stayed at 1.00x.
    supports_cuda_graph = True
    supports_mm_cluster_fix = False
    supports_triton_hooks = True
    # MUSA reports its own wording; none of the CUDA markers match it.
    fatal_error_markers = (
        "musa error: an illegal memory access",
        "out of memory",
        "device-side assert triggered",
        "misaligned address",
        "uncorrectable ecc",
    )
    # "unknown error" during capture is not recoverable on this runtime: the
    # fallback runs, and the process then dies in teardown with zero rows. Treat
    # it as poisoning so the operator is reported as capture-incompatible rather
    # than silently losing its measurement.
    poisoning_error_markers = ("musa error: unknown error",)
    # Keyed by FLAGGEMS_PERF_CURRENT_OP, which is the pytest FUNCTION name minus
    # `test_` -- `unique2`, `perf_sort` -- not the benchmark file name. Getting
    # that wrong is silent: the entry simply never matches.
    #
    # Observed on an S5000: capturing test_unique2 takes two recoverable
    # fallbacks and then raises "MUSA error: unknown error", ending the run at
    # rc=1 with no measured rows. Skipping the attempt costs that operator's
    # timing precision instead of its measurement. Revisit if a newer torch_musa
    # changes the behaviour.
    capture_deny_ops = ("unique2",)
    # Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while
    # the code under test dispatches through flag_gems.device. On MThreads that
    # turns a working operator into "not measured"; the plugin drops those
    # guards, pinned to one file and one exact skip reason each.
    compat_plugins = ("_musa_benchmark_compat_plugin",)
    # Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
    # a vendor CLI here; confirm the real listing tool before adding it.
    list_command = None

    # mthreads-gmi has no CSV query mode; --json is the only machine-readable
    # form it offers. Values arrive with their units glued on ("36C", "1888MHz",
    # "477.68W", "87%"), so the parser strips trailing non-numerics rather than
    # trusting a fixed suffix.
    telemetry_argv = (
        "mthreads-gmi", "-q",
        "-d", "TEMPERATURE,CLOCK,POWER,UTILIZATION",
        "--json",
        "-i", "{dev}",
    )

    def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
        import json
        import re

        def num(v):
            m = re.match(r"^\s*(-?[\d.]+)", str(v or ""))
            return m.group(1) if m else ""

        start = raw.find("{")
        if start < 0:
            return []
        try:
            data = json.loads(raw[start:])
        except Exception:
            return []
        gpus = data.get("GPU") or []
        if isinstance(gpus, dict):
            gpus = [gpus]
        out: List[Dict[str, str]] = []
        for g in gpus:
            clocks = g.get("Clocks") or {}
            power = g.get("Power Readings") or {}
            # 键名带尾随空格（"Power Draw "），别按字面写死
            draw = next((v for k, v in power.items()
                         if k.strip() == "Power Draw"), "")
            out.append({
                "index": str(g.get("Index", "")).strip(),
                "sm_clock_mhz": num(clocks.get("Graphics")),
                "mem_clock_mhz": num(clocks.get("Memory")),
                "temp_c": num((g.get("Temperature") or {}).get("GPU Current Temp")),
                "power_w": num(draw),
                "util_pct": num((g.get("Utilization") or {}).get("Gpu")),
            })
        return out