"""サーバの状況（GPU 一覧・実行中ジョブ）を表示用の文字列に整形する純関数群."""

from __future__ import annotations

from collections.abc import Sequence

from tias_judge_model.remote.gpu import (
    GpuInfo,
    GpuProcess,
    GpuSelection,
    GpuState,
    ServerProbe,
    display_user_name,
    gpu_users,
)
from tias_judge_model.remote.observe import format_elapsed

# 計算プロセスの一覧が取れなかった GPU の利用者の表示（run の警告と GPU 一覧で共通）
_UNKNOWN_USERS = "利用者不明（計算プロセスの一覧を取得できない）"


def _format_value(value: int | None, unit: str) -> str:
    return "?" if value is None else f"{value}{unit}"


def describe_gpu_users(gpu: GpuInfo, own_user: str) -> str:
    """GPU の利用者（計算プロセスの所有ユーザー）を表示用に 1 句で返す.

    Args:
        gpu: GPU の状況．
        own_user: 自分のユーザー名（ssh の接続ユーザー．一致するものは「自分」）．

    Returns:
        `利用者: 自分, alice` の形．計算プロセスの一覧が取れなければ `利用者不明`，
        計算プロセスが無ければ空文字．
    """
    if not gpu.processes_known:
        return _UNKNOWN_USERS
    users = gpu_users(gpu, own_user)
    return f"利用者: {', '.join(users)}" if users else ""


def describe_usage(gpu: GpuInfo, *, observed: bool) -> str:
    """GPU の使い方（型と使用率）を表示用に 1 句で返す.

    Args:
        gpu: GPU の状況．
        observed: 観測したか（False なら瞬間値だけを示し，型は付けない）．

    Returns:
        観測したときは `断続的に使用・平均 21%・最大 74%（0% 超 7/20）` の形．観測サンプルが
        無ければ `観測できない`．観測しないときは `使用率 29%（瞬間値）`．
    """
    usage = gpu.usage
    if not observed:
        return f"使用率 {_format_value(gpu.utilization_percent, '%')}（瞬間値）"
    if usage is None:
        return "観測できない（サンプル無し）"
    pattern = gpu.pattern.value if gpu.pattern is not None else "使い方不明"
    return (
        f"{pattern}・平均 {usage.mean_utilization_percent:.0f}%"
        f"・最大 {usage.max_utilization_percent}%"
        f"（0% 超 {usage.active_count}/{usage.sample_count}）"
    )


def describe_process(process: GpuProcess, own_user: str) -> str:
    """計算プロセス 1 件を表示用に 1 句で返す.

    Args:
        process: 計算プロセス．
        own_user: 自分のユーザー名（ssh の接続ユーザー．一致するものは「自分」）．

    Returns:
        `ando  経過 27.7 時間  SM 最大 74%` の形（取れない項目は省く．所有者が取れなければ
        `不明`）．
    """
    parts = [display_user_name(process.user, own_user)]
    if process.elapsed_sec is not None:
        parts.append(f"経過 {format_elapsed(process.elapsed_sec)}")
    if process.max_sm_percent is not None:
        parts.append(f"SM 最大 {process.max_sm_percent}%")
    return "  ".join(parts)


def _format_gpu_lines(gpu: GpuInfo, own_user: str, *, observed: bool) -> list[str]:
    index = "?" if gpu.index is None else str(gpu.index)
    memory = f"{_format_value(gpu.memory_used_mib, '')} / {_format_value(gpu.memory_total_mib, '')}"
    lines = [
        f"    GPU {index:>2}  {gpu.name}  メモリ {memory} MiB  [{gpu.state.value}]"
        f"  {describe_usage(gpu, observed=observed)}"
    ]
    if gpu.state is GpuState.UNKNOWN and gpu.usage is None and not gpu.processes:
        return lines
    if not gpu.processes_known:
        lines.append(f"        {_UNKNOWN_USERS}")
    else:
        lines.extend(f"        - {describe_process(p, own_user)}" for p in gpu.processes)
    return lines


def format_server_header(probe: ServerProbe) -> str:
    """サーバ 1 台の見出し行を返す.

    Args:
        probe: サーバの状況．

    Returns:
        `[名前] アドレス priority=N データ=... 観測 10 秒` の形の 1 行（接続不可なら理由を
        添える）．
    """
    server = probe.server
    data = "配置済み" if server.data_ready else "未配置"
    header = f"[{server.name}] {server.address}  priority={server.priority}  データ={data}"
    if probe.output is None:
        reason = f"（{probe.error}）" if probe.error else ""
        return f"{header}  接続不可{reason}"
    observe_sec = probe.output.observe_sec
    return f"{header}  {'観測 ' + str(observe_sec) + ' 秒' if observe_sec else '瞬間値のみ'}"


def format_server_report(probes: Sequence[ServerProbe], *, show_jobs: bool) -> str:
    """全サーバの GPU 一覧（状態と利用者．と実行中ジョブ）を表示用に整形する.

    Args:
        probes: 各サーバの状況（表示順）．
        show_jobs: 実行中ジョブも表示するか（status 用）．

    Returns:
        末尾に改行を含む複数行の文字列．
    """
    lines: list[str] = []
    for probe in probes:
        lines.append(format_server_header(probe))
        output = probe.output
        if output is None:
            continue
        if show_jobs:
            jobs = ", ".join(output.jobs) if output.jobs else "（なし）"
            lines.append(f"  実行中のジョブ: {jobs}")
            if output.other_sessions:
                lines.append(
                    f"  （このツール以外の tmux セッション {len(output.other_sessions)} 件）"
                )
        if not output.gpu_query_ok:
            lines.append("  GPU: nvidia-smi が失敗した（このサーバの GPU は候補から外す）")
        elif not output.gpus:
            lines.append("  GPU: なし")
        else:
            lines.append("  GPU:")
            for gpu in output.gpus:
                lines.extend(
                    _format_gpu_lines(gpu, probe.server.user, observed=output.observe_sec > 0)
                )
    return "\n".join(lines) + "\n"


def format_selection(selection: GpuSelection | None) -> str:
    """run が選ぶ GPU の説明を 1 行で返す.

    Args:
        selection: `select_gpu` の結果．

    Returns:
        表示用の 1 行（改行なし）．
    """
    if selection is None:
        return "run が選ぶ GPU: なし（未使用・空き（利用者あり）の GPU が無い）"
    text = (
        f"run が選ぶ GPU: {selection.server.name} の GPU {selection.gpu_index}"
        f"（{selection.state.value}）"
    )
    if selection.state is GpuState.SHARED and selection.gpu is not None:
        text += f"  {describe_gpu_users(selection.gpu, selection.server.user)}"
    return text


def format_shared_warning(selection: GpuSelection) -> str | None:
    """「空き（利用者あり）」の GPU を選んだときの警告文を返す（それ以外は None）.

    Args:
        selection: `select_gpu` の結果．

    Returns:
        警告文（改行なし）．
    """
    if selection.state is not GpuState.SHARED or selection.gpu is None:
        return None
    users = describe_gpu_users(selection.gpu, selection.server.user)
    return (
        f"{selection.server.name} の GPU {selection.gpu_index} には計算プロセスがある"
        f"（{users}）．他の計算とメモリ・計算時間を分け合う"
    )