from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias, TypeGuard

ShowFlag: TypeAlias = Literal[
    "showtemp",
    "showpower",
    "showsclk",
    "showmemuse",
    "showuse",
]


@dataclass(frozen=True)
class RenderColumn:
    label: str
    metric: str
    avg_label: str | None = None


@dataclass(frozen=True)
class ShowSpec:
    flag: ShowFlag
    metric_json_keys: dict[str, str]
    columns: tuple[RenderColumn, ...]
    cli_help: str
    hy_smi_flag: str | None = None


SHOW_SPECS: Final[tuple[ShowSpec, ...]] = (
    ShowSpec(
        flag="showtemp",
        metric_json_keys={"temp_c": "Temperature (Sensor core) (C)"},
        columns=(RenderColumn(label="Temp", metric="temp_c"),),
        cli_help="Display GPU core temperature.",
    ),
    ShowSpec(
        flag="showpower",
        metric_json_keys={"avg_pwr_w": "Average Graphics Package Power (W)"},
        columns=(RenderColumn(label="AvgPwr", metric="avg_pwr_w", avg_label="AvgPwr"),),
        cli_help="Display average GPU power.",
    ),
    ShowSpec(
        flag="showsclk",
        metric_json_keys={"sclk_mhz": "sclk clock speed"},
        columns=(RenderColumn(label="sclk", metric="sclk_mhz"),),
        cli_help="Display GPU sclk frequency.",
        hy_smi_flag="showhcuclocks",
    ),
    ShowSpec(
        flag="showmemuse",
        metric_json_keys={"vram_pct": "HCU memory use (%)"},
        columns=(RenderColumn(label="VRAM%", metric="vram_pct"),),
        cli_help="Display GPU VRAM usage.",
    ),
    ShowSpec(
        flag="showuse",
        metric_json_keys={"hcu_pct": "HCU use (%)"},
        columns=(RenderColumn(label="GPU%", metric="hcu_pct", avg_label="GPU%"),),
        cli_help="Display GPU utilization.",
    ),
)

SPEC_BY_FLAG: Final[dict[ShowFlag, ShowSpec]] = {spec.flag: spec for spec in SHOW_SPECS}

SUPPORTED_SHOW_FLAGS: Final[tuple[ShowFlag, ...]] = tuple(spec.flag for spec in SHOW_SPECS)

DEFAULT_SHOW_FLAGS: Final[tuple[ShowFlag, ...]] = SUPPORTED_SHOW_FLAGS

SHOW_FLAG_HELP: Final[dict[ShowFlag, str]] = {spec.flag: spec.cli_help for spec in SHOW_SPECS}

WAIT_IDLE_REQUIRED_SHOW_FLAGS: Final[tuple[ShowFlag, ...]] = ("showmemuse", "showuse")

JSON_KEY_BY_METRIC: Final[dict[str, str]] = {
    metric: json_key for spec in SHOW_SPECS for metric, json_key in spec.metric_json_keys.items()
}


def is_supported_show_flag(flag: str) -> TypeGuard[ShowFlag]:
    """Return True when the provided flag belongs to the supported show-flag set."""

    return flag in SPEC_BY_FLAG


def normalized_show_flags(show_flags: Iterable[str] | None) -> list[ShowFlag]:
    """Normalize user-selected show flags, preserving order and uniqueness."""

    if not show_flags:
        return list(DEFAULT_SHOW_FLAGS)
    output: list[ShowFlag] = []
    for flag in show_flags:
        spec = SPEC_BY_FLAG.get(flag)
        if spec is None:
            continue
        if spec.flag not in output:
            output.append(spec.flag)
    return output if output else list(DEFAULT_SHOW_FLAGS)


def hy_smi_args_for_show_flags(show_flags: Iterable[str], wait_idle: bool) -> list[str]:
    """Build hy-smi args with JSON output and requested fine-grained metrics."""

    ordered_flags = normalized_show_flags(show_flags)
    if wait_idle:
        # wait-idle relies on usage+memory metrics even when not displayed.
        for required in WAIT_IDLE_REQUIRED_SHOW_FLAGS:
            if required not in ordered_flags:
                ordered_flags.append(required)
    args = ["--json"]
    for flag in ordered_flags:
        spec = SPEC_BY_FLAG.get(flag)
        if spec:
            hy_flag = spec.hy_smi_flag or flag
            args.append(f"--{hy_flag}")
    return args


def render_columns_for_show_flags(show_flags: Iterable[str]) -> list[RenderColumn]:
    """Resolve display columns from ordered show flags."""

    columns: list[RenderColumn] = []
    for flag in normalized_show_flags(show_flags):
        columns.extend(SPEC_BY_FLAG[flag].columns)
    return columns
