"""テスト用に probe コマンド（SPEC_03「取得」）の出力文字列を組み立てる補助関数.

実サーバの nvidia-smi・ps を呼ばず，出力の文字列だけで判定を固定するために使う．
"""

from __future__ import annotations

from collections.abc import Sequence

from tias_judge_model.remote.commands import (
    PROBE_APPS_MARKER,
    PROBE_APPS_RC_PREFIX,
    PROBE_END_MARKER,
    PROBE_GPU_RC_PREFIX,
    PROBE_GPUS_MARKER,
    PROBE_JOB_PREFIX,
    PROBE_OTHER_PREFIX,
    PROBE_OWNERS_MARKER,
    PROBE_PMON_MARKER,
    PROBE_SAMPLES_MARKER,
    PROBE_SESSIONS_MARKER,
)

PMON_HEADER = [
    "# gpu         pid   type     sm    mem    enc    dec    jpg    ofa    command",
    "# Idx           #    C/G      %      %      %      %      %      %    name",
]


def gpu_line(index: int, uuid: str, used: int, util: int, *, total: int = 49140) -> str:
    """`nvidia-smi --query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu` の 1 行."""
    return f"{index}, {uuid}, NVIDIA RTX A6000, {used}, {total}, {util}"


def sample_lines(index: int, utils: Sequence[int], memory: int = 10) -> list[str]:
    """観測サンプル（`index, 使用率, 使用中メモリ`）の行."""
    return [f"{index}, {util}, {memory}" for util in utils]


def probe_text(
    gpus: Sequence[str],
    *,
    gpu_rc: int = 0,
    samples: Sequence[str] | None = None,
    pmon: Sequence[str] | None = None,
    apps: Sequence[str] | None = (),
    apps_rc: int = 0,
    owners: Sequence[str] = (),
    jobs: Sequence[str] = (),
    others: Sequence[str] = (),
    preamble: Sequence[str] = (),
) -> str:
    """probe の出力を組み立てる（samples・pmon が None なら節を出さない．apps が None なら
    計算プロセスの節ごと出さない）."""
    lines = [*preamble, PROBE_SESSIONS_MARKER]
    lines += [f"{PROBE_JOB_PREFIX}{name}" for name in jobs]
    lines += [f"{PROBE_OTHER_PREFIX}{name}" for name in others]
    lines.append(PROBE_GPUS_MARKER)
    lines += list(gpus)
    lines.append(f"{PROBE_GPU_RC_PREFIX}{gpu_rc}")
    if samples is not None:
        lines.append(PROBE_SAMPLES_MARKER)
        lines += list(samples)
    if pmon is not None:
        lines.append(PROBE_PMON_MARKER)
        lines += list(pmon)
    if apps is not None:
        lines.append(PROBE_APPS_MARKER)
        lines += list(apps)
        lines.append(f"{PROBE_APPS_RC_PREFIX}{apps_rc}")
        lines.append(PROBE_OWNERS_MARKER)
        lines += list(owners)
    lines.append(PROBE_END_MARKER)
    return "\n".join(lines) + "\n"