"""サーバの状況（tmux セッション・GPU・計算プロセス）の解析と，GPU の選択（純関数）.

ssh を呼ぶ部分は CLI 側にあり，ここは ssh の出力文字列を受け取って解釈するだけにする
（テストで ssh・nvidia-smi をモックせず，出力の文字列で固定できるようにするため）．
判定と選択の規則は SPEC_03「GPU の空き判定と選択」に従う．
"""

from __future__ import annotations

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

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,
)
from tias_judge_model.remote.config import GpuFreeThreshold, ServerConfig
from tias_judge_model.remote.observe import (
    GpuSample,
    GpuUsage,
    UsagePattern,
    classify_usage_pattern,
    parse_gpu_samples,
    parse_pmon,
    summarize_samples,
)

# nvidia-smi --query-gpu の列数（index,uuid,name,memory.used,memory.total,utilization.gpu）
_GPU_COLUMNS = 6
# nvidia-smi --query-compute-apps の列数（gpu_uuid,pid,used_memory）
_APP_COLUMNS = 3


class GpuState(StrEnum):
    """GPU の状態（選ぶ優先順は UNUSED → SHARED．BUSY・UNKNOWN は選ばない）."""

    UNUSED = "未使用"
    SHARED = "空き（利用者あり）"
    BUSY = "使用中"
    UNKNOWN = "状態不明"


# 自動選択で順に探す状態（未使用がサーバの優先順位より優先される）
SELECTABLE_STATES = (GpuState.UNUSED, GpuState.SHARED)


@dataclass(frozen=True)
class GpuProcess:
    """GPU 上の計算プロセス 1 件.

    Attributes:
        pid: プロセス ID．
        used_memory_mib: 使用メモリ（MiB）．取れなければ None．
        user: 所有ユーザー．ps で取れなければ None．
        elapsed_sec: 起動からの経過秒数．ps で取れなければ None．
        max_sm_percent: 観測中の SM 使用率の最大（pmon）．取れなければ None．
    """

    pid: int
    used_memory_mib: int | None
    user: str | None
    elapsed_sec: int | None = None
    max_sm_percent: int | None = None


@dataclass(frozen=True)
class ProcessOwner:
    """ps で取ったプロセスの所有者と経過時間.

    Attributes:
        user: 所有ユーザー．
        elapsed_sec: 起動からの経過秒数．取れなければ None．
    """

    user: str
    elapsed_sec: int | None


@dataclass(frozen=True)
class GpuInfo:
    """GPU 1 基の状況.

    Attributes:
        index: GPU 番号．解釈できなければ None．
        uuid: GPU の UUID（計算プロセスとの対応付けに使う）．解釈できない行では空文字．
        name: GPU の名前（解釈できない行では空文字）．
        memory_used_mib: 使用中メモリ（MiB．観測前の瞬間値）．解釈できなければ None．
        memory_total_mib: 総メモリ（MiB）．解釈できなければ None．
        utilization_percent: 使用率（%．観測前の瞬間値）．解釈できなければ None．
        state: 判定の結果．
        processes: この GPU 上の計算プロセス．
        processes_known: 計算プロセスの一覧が取れたか（False なら利用者不明）．
        usage: 観測の要約（観測しないときは瞬間値 1 つの要約）．サンプルが無ければ None．
        pattern: 使い方の型（表示専用）．決められなければ None．
    """

    index: int | None
    uuid: str
    name: str
    memory_used_mib: int | None
    memory_total_mib: int | None
    utilization_percent: int | None
    state: GpuState
    processes: tuple[GpuProcess, ...] = ()
    processes_known: bool = False
    usage: GpuUsage | None = None
    pattern: UsagePattern | None = None


@dataclass(frozen=True)
class ProbeOutput:
    """probe コマンド（`build_probe_command`）の出力を解釈した結果.

    Attributes:
        jobs: このツールのジョブ（ジョブログがある tmux セッション）の名前．
        other_sessions: それ以外の tmux セッションの名前．
        gpu_query_ok: GPU の状態の取得が成功したか．False なら GPU はすべて候補から外す．
        processes_query_ok: 計算プロセスの一覧の取得・解釈が成功したか．False なら閾値内の
            GPU はすべて「空き（利用者あり）」になる．
        gpus: GPU の一覧（GPU の状態の取得が失敗したときは空）．
        observe_sec: 観測した秒数（0 なら瞬間値のみで判定した）．
        pmon_ok: プロセス別の観測（pmon）が得られたか（False なら表示を省くだけ）．
    """

    jobs: tuple[str, ...]
    other_sessions: tuple[str, ...]
    gpu_query_ok: bool
    processes_query_ok: bool
    gpus: tuple[GpuInfo, ...]
    observe_sec: int = 0
    pmon_ok: bool = False

    @property
    def sessions(self) -> tuple[str, ...]:
        """すべての tmux セッション名（ジョブとそれ以外）."""
        return self.jobs + self.other_sessions


@dataclass(frozen=True)
class ServerProbe:
    """サーバ 1 台の状況.

    Attributes:
        server: サーバの設定．
        output: probe の結果．接続できなかったときは None．
        error: 接続できなかった理由（表示用．ssh の標準エラーの最終行等）．
    """

    server: ServerConfig
    output: ProbeOutput | None
    error: str = ""

    @property
    def reachable(self) -> bool:
        """接続できたか."""
        return self.output is not None


@dataclass(frozen=True)
class GpuSelection:
    """run が使う GPU.

    Attributes:
        server: 選んだサーバ．
        gpu_index: 選んだ GPU 番号．
        state: 選んだ時点の GPU の状態．
        gpu: 選んだ GPU の状況（利用者の警告表示用．状態を取れなかったときは None）．
        forced: 使用中・状態不明の GPU を `--force` で選んだか（警告を出す）．
    """

    server: ServerConfig
    gpu_index: int
    state: GpuState
    gpu: GpuInfo | None = None
    forced: bool = False


def classify_gpu(
    memory_used_mib: int,
    utilization_percent: int,
    threshold: GpuFreeThreshold,
    process_count: int | None,
) -> GpuState:
    """閾値と計算プロセスの数から GPU の状態を判定する.

    観測するときは，使用中メモリ・使用率に**観測中の最大値**を渡す（一度でも閾値を超えた
    GPU を使用中にするため）．

    Args:
        memory_used_mib: 使用中メモリ（MiB．観測中の最大）．
        utilization_percent: 使用率（%．観測中の最大）．
        threshold: 空き判定の閾値．
        process_count: 計算プロセスの数（所有者を問わない）．一覧が取れなければ None．

    Returns:
        閾値を満たさなければ `BUSY`．閾値内で計算プロセスが 0 なら `UNUSED`，
        1 つ以上または不明（None）なら `SHARED`（未使用と誤認しないため）．
    """
    within = (
        memory_used_mib < threshold.max_memory_used_mib
        and utilization_percent < threshold.max_utilization_percent
    )
    if not within:
        return GpuState.BUSY
    return GpuState.UNUSED if process_count == 0 else GpuState.SHARED


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


def _is_valid_uuid(text: str) -> bool:
    """nvidia-smi の UUID として使える値か（空・`[N/A]` 等を除く）."""
    return bool(text) and not text.startswith("[") and " " not in text


def parse_process_owners(text: str) -> dict[int, ProcessOwner]:
    """probe の所有者の節（`<pid> <ユーザー> <経過秒>` の行）を解釈する.

    ユーザー名が無い行（ps で所有者を取れなかったプロセス）と解釈できない行は含めない
    （所有者・経過時間は表示用で，判定には使わないため）．

    Args:
        text: 所有者の節の出力．

    Returns:
        pid から所有者・経過秒数への対応．
    """
    owners: dict[int, ProcessOwner] = {}
    for line in text.splitlines():
        tokens = line.split()
        if len(tokens) < 2:
            continue
        pid = _parse_int(tokens[0])
        if pid is None:
            continue
        elapsed = _parse_int(tokens[2]) if len(tokens) >= 3 else None
        owners[pid] = ProcessOwner(tokens[1], elapsed)
    return owners


def parse_compute_apps(
    text: str,
    owners: Mapping[int, ProcessOwner],
    sm_peaks: Mapping[int, int] | None = None,
) -> dict[str, tuple[GpuProcess, ...]] | None:
    """`nvidia-smi --query-compute-apps=gpu_uuid,pid,used_memory` の出力を GPU ごとにまとめる.

    Args:
        text: nvidia-smi の標準出力（空なら計算プロセスは 0）．
        owners: pid から所有者・経過秒数への対応（`parse_process_owners`）．
        sm_peaks: pid から観測中の SM 使用率の最大への対応（`parse_pmon`）．無ければ None．

    Returns:
        GPU の UUID から計算プロセスへの対応．1 行でも解釈できなければ None
        （一覧を信用せず利用者不明として扱うため）．
    """
    apps: dict[str, list[GpuProcess]] = {}
    for line in text.splitlines():
        if not line.strip():
            continue
        columns = [column.strip() for column in line.split(",")]
        if len(columns) != _APP_COLUMNS or not _is_valid_uuid(columns[0]):
            return None
        pid = _parse_int(columns[1])
        if pid is None:
            return None
        owner = owners.get(pid)
        process = GpuProcess(
            pid=pid,
            used_memory_mib=_parse_int(columns[2]),
            user=owner.user if owner else None,
            elapsed_sec=owner.elapsed_sec if owner else None,
            max_sm_percent=sm_peaks.get(pid) if sm_peaks is not None else None,
        )
        apps.setdefault(columns[0], []).append(process)
    return {uuid: tuple(processes) for uuid, processes in apps.items()}


def parse_nvidia_smi_line(
    line: str,
    threshold: GpuFreeThreshold,
    processes: Mapping[str, Sequence[GpuProcess]] | None,
    samples: Sequence[GpuSample] | None = None,
) -> GpuInfo:
    """`nvidia-smi --query-gpu`（`--format=csv,noheader,nounits`）の 1 行を解釈して判定する.

    列数不足・数値でない値（`[N/A]` 等）の行は `UNKNOWN` とする．GPU 名にカンマを含む場合に
    備え，先頭 2 列を index・uuid，末尾 3 列を数値として，その間を名前とする．

    観測したとき（`samples` が None でない）は，観測中の最大値（とこの行の瞬間値の大きい方）で
    閾値を判定し，サンプルが 1 つも無ければ `UNKNOWN` とする．観測しないとき（None）は
    この行の瞬間値で判定する．

    Args:
        line: nvidia-smi の出力 1 行．
        threshold: 空き判定の閾値．
        processes: GPU の UUID から計算プロセスへの対応（`parse_compute_apps`）．
            計算プロセスの一覧が取れなかったときは None（閾値内なら `SHARED` になる）．
        samples: この GPU の観測サンプル．観測しないときは None．

    Returns:
        GPU の状況．
    """
    columns = [column.strip() for column in line.split(",")]
    if len(columns) < _GPU_COLUMNS:
        index = _parse_int(columns[0]) if columns else None
        return GpuInfo(index, "", "", None, None, None, GpuState.UNKNOWN)
    index = _parse_int(columns[0])
    uuid = columns[1]
    name = ", ".join(columns[2:-3])
    used, total, util = (_parse_int(value) for value in columns[-3:])
    # UUID が取れない GPU は計算プロセスと対応付けられないので，利用者不明として扱う
    known = processes is not None and _is_valid_uuid(uuid)
    gpu_processes = tuple(processes.get(uuid, ())) if known and processes is not None else ()
    count = len(gpu_processes) if known else None

    if index is None or used is None or total is None or util is None:
        return GpuInfo(index, uuid, name, used, total, util, GpuState.UNKNOWN, gpu_processes, known)
    observed = samples is not None
    usage = summarize_samples(samples if samples is not None else [GpuSample(util, used)])
    if usage is None:
        # 観測したのにサンプルが 1 つも得られない GPU は使わない
        state = GpuState.UNKNOWN
        pattern = None
    else:
        max_used = max(usage.max_memory_used_mib, used) if observed else used
        max_util = max(usage.max_utilization_percent, util) if observed else util
        state = classify_gpu(max_used, max_util, threshold, count)
        pattern = classify_usage_pattern(usage, count)
    return GpuInfo(
        index, uuid, name, used, total, util, state, gpu_processes, known, usage, pattern
    )


def parse_nvidia_smi_output(
    text: str,
    threshold: GpuFreeThreshold,
    processes: Mapping[str, Sequence[GpuProcess]] | None,
    samples: Mapping[int, Sequence[GpuSample]] | None = None,
) -> list[GpuInfo]:
    """`nvidia-smi --query-gpu` の出力全体を 1 行 1 GPU として解釈する（空行は無視する）.

    Args:
        text: nvidia-smi の標準出力．
        threshold: 空き判定の閾値．
        processes: GPU の UUID から計算プロセスへの対応．一覧が取れなかったときは None．
        samples: GPU 番号から観測サンプルへの対応（`parse_gpu_samples`）．観測しないときは
            None（瞬間値で判定する）．観測したのに該当する GPU のサンプルが無ければ状態不明．

    Returns:
        出力順の GPU の一覧．
    """
    gpus: list[GpuInfo] = []
    for line in text.splitlines():
        if not line.strip():
            continue
        gpu_samples: Sequence[GpuSample] | None = None
        if samples is not None:
            index = _parse_int(line.split(",", 1)[0])
            gpu_samples = samples.get(index, []) if index is not None else []
        gpus.append(parse_nvidia_smi_line(line, threshold, processes, gpu_samples))
    return gpus


def parse_probe_output(text: str, threshold: GpuFreeThreshold) -> ProbeOutput:
    """probe コマンド（`build_probe_command`）の標準出力を解釈する.

    目印より前の行（ログインシェルが出す挨拶等）は無視する．計算プロセスの節が欠けている・
    終了コードが 0 でない・解釈できない場合は，計算プロセスの一覧を「取得できない」とする．
    `threshold.observe_sec > 0` なら観測したものとして判定し，観測の節が欠けていれば
    （サーバで一時ディレクトリを作れなかった等）各 GPU は状態不明になる．

    Args:
        text: probe コマンドの標準出力．
        threshold: 空き判定の閾値と観測の設定（probe コマンドの組み立てに使ったもの）．

    Returns:
        解釈した結果．

    Raises:
        ValueError: 必須の目印（セッション・GPU・GPU の終了コード・終端）がそろっていない
            （出力が途中で切れた等）場合．
    """
    section: str | None = None
    jobs: list[str] = []
    others: list[str] = []
    lines: dict[str, list[str]] = {
        "gpus": [],
        "samples": [],
        "pmon": [],
        "apps": [],
        "owners": [],
    }
    gpu_rc: int | None = None
    apps_rc: int | None = None
    seen_sessions = ended = False
    for line in text.splitlines():
        if line == PROBE_SESSIONS_MARKER:
            section, seen_sessions = "sessions", True
        elif line == PROBE_GPUS_MARKER:
            section = "gpus"
        elif line.startswith(PROBE_GPU_RC_PREFIX) and section == "gpus":
            gpu_rc = _parse_int(line[len(PROBE_GPU_RC_PREFIX) :])
            section = None
        elif line == PROBE_SAMPLES_MARKER:
            section = "samples"
        elif line == PROBE_PMON_MARKER:
            section = "pmon"
        elif line == PROBE_APPS_MARKER:
            section = "apps"
        elif line.startswith(PROBE_APPS_RC_PREFIX) and section == "apps":
            apps_rc = _parse_int(line[len(PROBE_APPS_RC_PREFIX) :])
            section = None
        elif line == PROBE_OWNERS_MARKER:
            section = "owners"
        elif line == PROBE_END_MARKER:
            section, ended = None, True
        elif section == "sessions":
            if line.startswith(PROBE_JOB_PREFIX):
                jobs.append(line[len(PROBE_JOB_PREFIX) :])
            elif line.startswith(PROBE_OTHER_PREFIX):
                others.append(line[len(PROBE_OTHER_PREFIX) :])
        elif section in lines:
            lines[section].append(line)
    if not (seen_sessions and ended) or gpu_rc is None:
        raise ValueError("サーバの状況の出力を解釈できない（目印が欠けている）")

    texts = {name: "\n".join(section_lines) for name, section_lines in lines.items()}
    gpu_ok = gpu_rc == 0
    observed = threshold.observe_sec > 0
    sm_peaks = parse_pmon(texts["pmon"]) if observed else None
    processes = None
    if gpu_ok and apps_rc == 0:
        owners = parse_process_owners(texts["owners"])
        processes = parse_compute_apps(texts["apps"], owners, sm_peaks)
    samples = parse_gpu_samples(texts["samples"]) if observed else None
    gpus = (
        tuple(parse_nvidia_smi_output(texts["gpus"], threshold, processes, samples))
        if gpu_ok
        else ()
    )
    return ProbeOutput(
        jobs=tuple(jobs),
        other_sessions=tuple(others),
        gpu_query_ok=gpu_ok,
        processes_query_ok=processes is not None,
        gpus=gpus,
        observe_sec=threshold.observe_sec if observed else 0,
        pmon_ok=sm_peaks is not None,
    )


def _find_probe(probes: Sequence[ServerProbe], name: str) -> ServerProbe:
    """名前（大文字小文字を区別しない）でサーバの状況を探す."""
    for probe in probes:
        if probe.server.name.lower() == name.lower():
            return probe
    known = ", ".join(probe.server.name for probe in probes)
    raise ValueError(f"サーバ {name!r} は設定に無い（設定済み: {known}）")


def _select_explicit(probe: ServerProbe, gpu_index: int, *, force: bool) -> GpuSelection:
    """`--server --gpu` で指定された GPU を検証して返す."""
    output = probe.output
    if output is None:
        raise RuntimeError("接続できないサーバの GPU を検証しようとした")
    gpu = next((g for g in output.gpus if g.index == gpu_index), None)
    if gpu is None:
        # 一覧を正しく取れていて該当が無いなら，存在しない GPU の指定である
        if output.gpu_query_ok and all(g.index is not None for g in output.gpus):
            raise ValueError(f"{probe.server.name} に GPU {gpu_index} は無い")
        state = GpuState.UNKNOWN
    else:
        state = gpu.state
    if state in SELECTABLE_STATES:
        return GpuSelection(probe.server, gpu_index, state, gpu)
    if not force:
        raise ValueError(
            f"{probe.server.name} の GPU {gpu_index} は{state}である"
            "（使うなら --force を付ける．他のユーザーが使用中の可能性がある）"
        )
    return GpuSelection(probe.server, gpu_index, state, gpu, forced=True)


def check_selection_options(*, server_name: str | None, gpu_index: int | None, force: bool) -> None:
    """GPU 指定の組み合わせを検証する（サーバへ問い合わせる前に弾けるよう独立させている）.

    Args:
        server_name: `--server` の値．
        gpu_index: `--gpu` の値．
        force: `--force` の有無．

    Raises:
        ValueError: `--gpu` だけの指定，`--gpu` 無しの `--force`，負の GPU 番号の場合．
    """
    if gpu_index is not None and server_name is None:
        raise ValueError("--gpu を指定するときは --server も指定する（どのサーバか曖昧になるため）")
    if force and gpu_index is None:
        raise ValueError("--force は --server と --gpu で GPU を指定したときだけ使える")
    if gpu_index is not None and gpu_index < 0:
        raise ValueError(f"GPU 番号は 0 以上: {gpu_index}")


def select_gpu(
    probes: Sequence[ServerProbe],
    *,
    server_name: str | None = None,
    gpu_index: int | None = None,
    force: bool = False,
    include_not_ready: bool = False,
) -> GpuSelection | None:
    """run が使う GPU を選ぶ（SPEC_03「GPU の空き判定と選択」）.

    自動選択では，候補サーバ（`data_ready` のもの．`include_not_ready` なら全サーバ．
    `server_name` を指定したらそのサーバのみ）全体を通して，まず「未使用」を priority 昇順
    （同値は `probes` の順）→ GPU 番号の昇順で探し，無ければ同じ順で「空き（利用者あり）」を
    探す（未使用がサーバの優先順位より優先される）．接続できないサーバと GPU の状態の取得に
    失敗したサーバは候補から外す．

    Args:
        probes: 各サーバの状況（設定の記載順で渡す．同じ priority の順位に使う）．
        server_name: `--server` の値．
        gpu_index: `--gpu` の値（`server_name` が必須）．
        force: 指定した GPU が使用中・状態不明でも選ぶ（`gpu_index` 指定時のみ有効）．
        include_not_ready: `data_ready: false` のサーバも候補にする（`--no-data`）．

    Returns:
        選んだ GPU．自動選択で未使用・空き（利用者あり）のどちらも無ければ None．

    Raises:
        ValueError: 指定の組み合わせが不正（`--gpu` だけ・`--force` だけ），指定したサーバが
            設定に無い・データ未配置・接続不可，指定した GPU が無い・使用中・状態不明
            （`force` 無し）の場合．
    """
    check_selection_options(server_name=server_name, gpu_index=gpu_index, force=force)
    if server_name is not None:
        probe = _find_probe(probes, server_name)
        if not probe.server.data_ready and not include_not_ready:
            raise ValueError(
                f"{probe.server.name} は臨床データが未配置（data_ready: false）である"
                "（データを使わないジョブなら --no-data を付ける）"
            )
        if not probe.reachable:
            raise ValueError(f"{probe.server.name} に接続できない（{probe.error}）")
        if gpu_index is not None:
            return _select_explicit(probe, gpu_index, force=force)
        candidates = [probe]
    else:
        candidates = sorted(
            (p for p in probes if p.reachable and (include_not_ready or p.server.data_ready)),
            key=lambda p: p.server.priority,
        )

    for wanted in SELECTABLE_STATES:
        for probe in candidates:
            output = probe.output
            if output is None or not output.gpu_query_ok:
                continue
            for gpu in sorted(
                (g for g in output.gpus if g.index is not None), key=lambda g: g.index or 0
            ):
                if gpu.state is wanted and gpu.index is not None:
                    return GpuSelection(probe.server, gpu.index, gpu.state, gpu)
    return None


def display_user_name(user: str | None, own_user: str) -> str:
    """計算プロセスの所有ユーザーを表示用の名前にする.

    Args:
        user: 所有ユーザー．ps で取れなければ None．
        own_user: 自分のユーザー名（ssh の接続ユーザー）．

    Returns:
        自分なら「自分」，取れなければ「不明」，それ以外はユーザー名．
    """
    if user is None:
        return "不明"
    if user == own_user:
        return "自分"
    return user


def gpu_users(gpu: GpuInfo, own_user: str) -> list[str]:
    """GPU の計算プロセスの所有ユーザーを表示用に並べる（重複は 1 つにまとめ，出現順）.

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

    Returns:
        ユーザー名の一覧（自分は「自分」，所有者を取れなかったものは「不明」）．
    """
    names: list[str] = []
    for process in gpu.processes:
        name = display_user_name(process.user, own_user)
        if name not in names:
            names.append(name)
    return names


def servers_with_session(
    probes: Sequence[ServerProbe], name: str, *, jobs_only: bool
) -> list[ServerProbe]:
    """指定した名前の tmux セッションがあるサーバを返す.

    Args:
        probes: 各サーバの状況（接続できなかったサーバは対象外）．
        name: セッション（ジョブ）名．
        jobs_only: True ならこのツールのジョブ（ジョブログがあるもの）だけを見る．

    Returns:
        該当するサーバの状況（`probes` の順）．
    """
    found: list[ServerProbe] = []
    for probe in probes:
        if probe.output is None:
            continue
        names = probe.output.jobs if jobs_only else probe.output.sessions
        if name in names:
            found.append(probe)
    return found