"""GPU の観測結果（サンプリング・pmon）の解析と，使い方の型の判定（純関数）.

観測は probe コマンドがサーバ側で行い（`build_probe_command`），ここはその出力文字列を
解釈するだけにする．使い方の型は表示専用で，空き判定には使わない（SPEC_03「使い方の型」）．
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from enum import StrEnum

# 0% を超えたサンプルがこの割合以上なら「常時使用」
CONSTANT_USE_RATIO = 0.8
# nvidia-smi --query-gpu=index,utilization.gpu,memory.used の列数
_SAMPLE_COLUMNS = 3


class UsagePattern(StrEnum):
    """観測中の使用率から付ける使い方の型（表示専用）."""

    CONSTANT = "常時使用"
    INTERMITTENT = "断続的に使用"
    MEMORY_ONLY = "メモリ確保のみ"
    IDLE = "未使用"


@dataclass(frozen=True)
class GpuSample:
    """GPU 1 基の 1 回分のサンプル.

    Attributes:
        utilization_percent: 使用率（%）．
        memory_used_mib: 使用中メモリ（MiB）．
    """

    utilization_percent: int
    memory_used_mib: int


@dataclass(frozen=True)
class GpuUsage:
    """観測中のサンプルの要約.

    Attributes:
        sample_count: サンプル数（1 以上）．
        active_count: 使用率が 0% を超えたサンプル数．
        mean_utilization_percent: 使用率の平均（%）．
        max_utilization_percent: 使用率の最大（%）．
        max_memory_used_mib: 使用中メモリの最大（MiB）．
    """

    sample_count: int
    active_count: int
    mean_utilization_percent: float
    max_utilization_percent: int
    max_memory_used_mib: int


def _parse_int(text: str) -> int | None:
    """整数として解釈できれば返す（`[N/A]`・`-` 等は None）."""
    try:
        return int(text.strip())
    except ValueError:
        return None


def parse_gpu_samples(text: str) -> dict[int, list[GpuSample]]:
    """観測のサンプリング出力（`index, 使用率, 使用中メモリ` の CSV）を GPU ごとにまとめる.

    `timeout` で打ち切った最後の行が途中で切れていることがあるため，解釈できない行は
    読み飛ばす（そのぶんサンプルが減るだけで，サンプルが 0 件の GPU は状態不明になる）．

    Args:
        text: サンプリングの出力．

    Returns:
        GPU 番号からサンプル（出力順）への対応．
    """
    samples: dict[int, list[GpuSample]] = {}
    for line in text.splitlines():
        columns = line.split(",")
        if len(columns) != _SAMPLE_COLUMNS:
            continue
        index, util, memory = (_parse_int(column) for column in columns)
        if index is None or util is None or memory is None:
            continue
        samples.setdefault(index, []).append(GpuSample(util, memory))
    return samples


def summarize_samples(samples: Sequence[GpuSample]) -> GpuUsage | None:
    """サンプルを要約する.

    Args:
        samples: GPU 1 基のサンプル．

    Returns:
        要約．サンプルが無ければ None．
    """
    if not samples:
        return None
    utils = [sample.utilization_percent for sample in samples]
    return GpuUsage(
        sample_count=len(samples),
        active_count=sum(1 for util in utils if util > 0),
        mean_utilization_percent=sum(utils) / len(utils),
        max_utilization_percent=max(utils),
        max_memory_used_mib=max(sample.memory_used_mib for sample in samples),
    )


def classify_usage_pattern(usage: GpuUsage, process_count: int | None) -> UsagePattern | None:
    """観測中の使用率から使い方の型を付ける（表示専用）.

    Args:
        usage: 観測の要約．
        process_count: 計算プロセスの数．一覧が取れなければ None．

    Returns:
        使い方の型．0% を超えたサンプルが無く，計算プロセスの有無も分からなければ None．
    """
    if usage.active_count > 0:
        if usage.active_count >= CONSTANT_USE_RATIO * usage.sample_count:
            return UsagePattern.CONSTANT
        return UsagePattern.INTERMITTENT
    if process_count is None:
        return None
    return UsagePattern.MEMORY_ONLY if process_count > 0 else UsagePattern.IDLE


def parse_pmon(text: str) -> dict[int, int] | None:
    """`nvidia-smi pmon -s u` の出力から，計算プロセスごとの SM 使用率の最大を求める.

    `#` で始まる行はヘッダ（1 行目が列名 `# gpu pid type sm ...`，2 行目が単位）．列の位置は
    ヘッダの列名から決める．種別に `C`（計算）を含む行だけを数え，`G`（描画）だけの行は除く．
    値が `-` の行（その回はプロセスが無い・値が取れない）は読み飛ばす．

    Args:
        text: pmon の出力．

    Returns:
        pid から SM 使用率の最大（%）への対応．列名のヘッダが無い（pmon が失敗・未対応）
        ときは None（表示を省くだけで，判定には影響させない）．
    """
    columns: list[str] | None = None
    peaks: dict[int, int] = {}
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        if stripped.startswith("#"):
            names = stripped.lstrip("#").split()
            if columns is None and {"pid", "type", "sm"} <= set(names):
                columns = names
            continue
        if columns is None:
            continue
        values = stripped.split()
        if len(values) < len(columns) - 1:
            continue
        row = dict(zip(columns, values, strict=False))
        pid = _parse_int(row.get("pid", "-"))
        sm = _parse_int(row.get("sm", "-"))
        if pid is None or sm is None or "C" not in row.get("type", ""):
            continue
        peaks[pid] = max(peaks.get(pid, 0), sm)
    return None if columns is None else peaks


def format_elapsed(seconds: int) -> str:
    """経過秒数を読みやすい単位で返す（例: `45 秒`・`12.5 分`・`27.7 時間`・`3.2 日`）.

    48 時間までは時間で示す（日単位にすると 1 日前後の違いが読み取りにくいため）．

    Args:
        seconds: 経過秒数（0 以上）．

    Returns:
        表示用の文字列．
    """
    if seconds < 60:
        return f"{seconds} 秒"
    if seconds < 3600:
        return f"{seconds / 60:.1f} 分"
    if seconds < 48 * 3600:
        return f"{seconds / 3600:.1f} 時間"
    return f"{seconds / 86400:.1f} 日"