"""tias-remote: 研究室の GPU サーバでの実行を手元から操作する CLI（SPEC_03）.

手元（Windows の PowerShell / Git Bash）から 1 コマンドで以下を行う．サーバは
`configs/remote.yaml` の一覧（Newton・Archimedes 等）で，空いている GPU を自動で選ぶ．

- gpus   : 全サーバの GPU 一覧と，run が選ぶ GPU を表示する
- status : 全サーバの実行中ジョブと GPU の状況を表示する
- push   : コードをサーバへ送る（tar over ssh．rsync 不要）
- run    : 空き GPU を選び，そのサーバへコードを送って tmux でジョブを起動する
- logs   : ジョブのログを表示する（ジョブのあるサーバを自動で探す）
- kill   : ジョブを止める（ジョブのあるサーバを自動で探す）
- pull   : 実験結果を回収ホワイトリストに従って持ち帰る

ssh を呼ぶ副作用はこのモジュールに集め，コマンド文字列の組み立て・出力の解析・GPU の選択・
回収ホワイトリストの判定は `tias_judge_model.remote` の純関数に任せる．
"""

from __future__ import annotations

import argparse
import codecs
import contextlib
import logging
import os
import socket
import subprocess
import sys
import tarfile
import tempfile
import threading
import time
from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO

from pydantic import ValidationError

from tias_judge_model.remote.commands import (
    EXIT_NO_SESSION,
    EXIT_NOT_FOUND,
    EXIT_REMOTE_DIR_MISSING,
    EXIT_SESSION_EXISTS,
    EXIT_UNSAFE_PATH,
    REMOTE_LOG_DIR,
    REMOTE_RUNS_DIR,
    build_job_env,
    build_job_script,
    build_kill_command,
    build_log_exists_command,
    build_logs_command,
    build_probe_command,
    build_pull_list_command,
    build_pull_tar_command,
    build_push_command,
    build_run_command,
    build_ssh_argv,
    format_git_commit,
    job_log_path,
    validate_job_name,
)
from tias_judge_model.remote.config import (
    ALLOW_SELF_ENV_VAR,
    DEFAULT_CONFIG_PATH,
    JUMP_ENV_VAR,
    SSH_ENV_VAR,
    GpuFreeThreshold,
    RemoteConfig,
    ServerConfig,
    detect_self_server,
    is_self_check_disabled,
    load_remote_config,
    resolve_jump_host,
    resolve_ssh_executable,
)
from tias_judge_model.remote.files import PushEntry, collect_push_entries, find_project_root
from tias_judge_model.remote.gpu import (
    ServerProbe,
    check_cpu_options,
    check_selection_options,
    parse_probe_output,
    select_cpu_server,
    select_gpu,
    servers_with_session,
)
from tias_judge_model.remote.report import (
    format_selection,
    format_server_report,
    format_shared_warning,
)
from tias_judge_model.remote.transfer import copy_stream, extract_pull_archive, write_tar_stream
from tias_judge_model.remote.whitelist import parse_find_listing, plan_pull, validate_run_id
from tias_judge_model.schema import DATA_ROOT_ENV_VAR

logger = logging.getLogger(__name__)

# ssh 自身の失敗（接続できない・認証失敗等）の終了コード
_SSH_FAILURE_RC = 255
# 応答待ちを打ち切ったことを表す内部の終了コード（ssh の終了コードと衝突しない負値）
_TIMEOUT_RC = -1
# GNU tar の終了コード 1 は「読み取り中にファイルが変化した」等の警告
_TAR_WARNING_RC = 1
# 短い問い合わせ（probe・ログの有無・一覧）の応答待ちの上限は接続タイムアウトにこれを足す．
# 踏み台経由では ConnectTimeout が踏み台側の接続に効かないため，全体の上限も設ける
_QUERY_TIMEOUT_MARGIN_SEC = 30
_MAX_STDERR_CHARS = 160
_HOST_KEY_FAILURE = b"Host key verification failed"
# 手元の git（rev-parse・status）の応答待ちの上限
_GIT_TIMEOUT_SEC = 30

_EPILOG = """\
使用例:
  uv run tias-remote gpus                                   # 全サーバの GPU と run が選ぶ GPU
  uv run tias-remote status                                 # 実行中ジョブと GPU
  uv run tias-remote push                                   # 接続できる全サーバへコードを送る
  uv run tias-remote run --name exp1 -- pytest              # 空き GPU を自動で選んで起動
  uv run tias-remote run --name exp2 --server Newton --gpu 1 -- python -m ...   # GPU を指定
  uv run tias-remote run --name survey --cpu -- tias-dataset survey --config ...   # GPU なし
  uv run tias-remote logs exp1 -f                           # ログを追従
  uv run tias-remote kill exp1                              # ジョブを止める
  uv run tias-remote pull 20260927-120000_newton_abc1234    # 回収ホワイトリストのみ持ち帰る

環境変数:
  TIAS_REMOTE_JUMP        踏み台を上書き（空文字なら直接接続．研究室 LAN 内にいるとき）
  TIAS_REMOTE_SSH         手元の ssh 実行ファイル
  TIAS_REMOTE_ALLOW_SELF  1 でサーバ上実行の検出を無効化
"""


@dataclass(frozen=True)
class _Context:
    """サブコマンドが使う実行環境.

    Attributes:
        config: サーバ設定．
        ssh: 手元の ssh 実行ファイル．
        jump_host: 踏み台（None なら直接接続）．
        project_root: 手元のプロジェクトルート．
    """

    config: RemoteConfig
    ssh: str
    jump_host: str | None
    project_root: Path

    @property
    def query_timeout_sec(self) -> int:
        """短い問い合わせの応答待ちの上限（秒）."""
        return self.config.connect_timeout_sec + _QUERY_TIMEOUT_MARGIN_SEC


@dataclass(frozen=True)
class _SshResult:
    """ssh の実行結果.

    Attributes:
        rc: 終了コード（応答待ちを打ち切ったら `_TIMEOUT_RC`）．
        stdout: 標準出力．
        stderr: 標準エラー．
    """

    rc: int
    stdout: bytes
    stderr: bytes


class _CliFormatter(logging.Formatter):
    """警告・エラーに日本語の接頭辞を付ける（CLI の利用者向けの表示）."""

    _PREFIXES = {logging.WARNING: "警告: ", logging.ERROR: "エラー: ", logging.CRITICAL: "エラー: "}

    def format(self, record: logging.LogRecord) -> str:
        return self._PREFIXES.get(record.levelno, "") + super().format(record)


def _configure_logging() -> None:
    """標準エラーへ日本語の接頭辞付きでログを出す（既に設定済みなら何もしない）."""
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(_CliFormatter("%(message)s"))
    logging.basicConfig(level=logging.INFO, handlers=[handler])


def _write_stdout(text: str) -> None:
    """標準出力へ書く（出力先の文字コードで表せない文字は置換する）.

    リモートの出力は UTF-8 で届くが，リダイレクト時の Windows の標準出力は cp932 のことがある．
    """
    try:
        sys.stdout.write(text)
    except UnicodeEncodeError:
        encoding = sys.stdout.encoding or "utf-8"
        sys.stdout.write(text.encode(encoding, errors="replace").decode(encoding, errors="replace"))
    sys.stdout.flush()


def _format_size(num_bytes: float) -> str:
    """バイト数を読みやすい単位で返す."""
    for unit in ("B", "KB", "MB"):
        if num_bytes < 1024:
            return f"{num_bytes:.0f} {unit}" if unit == "B" else f"{num_bytes:.1f} {unit}"
        num_bytes /= 1024
    return f"{num_bytes:.1f} GB"


def _last_stderr_line(stderr: bytes) -> str:
    """ssh の標準エラーの最後の空でない行を返す（接続不可の理由の表示用）."""
    lines = [line.strip() for line in stderr.decode("utf-8", errors="replace").splitlines()]
    lines = [line for line in lines if line]
    return lines[-1][:_MAX_STDERR_CHARS] if lines else ""


def _describe_failure(result: _SshResult) -> str:
    """ssh が失敗した理由を短く返す."""
    if result.rc == _TIMEOUT_RC:
        return "応答がタイムアウトした"
    detail = _last_stderr_line(result.stderr) if result.rc == _SSH_FAILURE_RC else ""
    return f"終了コード {result.rc}" + (f": {detail}" if detail else "")


# ---------------------------------------------------------------------------
# ssh の実行（副作用）
# ---------------------------------------------------------------------------


def _ssh_argv(
    ctx: _Context, server: ServerConfig, command: str, *, stdin: bool = False
) -> list[str]:
    """サーバでリモートコマンドを実行する ssh の引数リストを返す."""
    return build_ssh_argv(
        ctx.ssh,
        user=server.user,
        address=server.address,
        jump_host=ctx.jump_host,
        connect_timeout_sec=ctx.config.connect_timeout_sec,
        remote_command=command,
        stdin=stdin,
    )


def _popen(argv: list[str], **kwargs: object) -> subprocess.Popen[bytes]:
    """ssh を起動する（見つからなければ対処を添えて ValueError にする）."""
    try:
        return subprocess.Popen(argv, **kwargs)  # noqa: S603
    except FileNotFoundError as exc:
        raise ValueError(
            f"ssh が見つからない: {argv[0]}（環境変数 {SSH_ENV_VAR} で実行ファイルを指定できる）"
        ) from exc


def _run_ssh(
    argv: list[str], *, input_bytes: bytes | None = None, timeout: float | None = None
) -> _SshResult:
    """ssh を実行して終了を待ち，標準出力・標準エラーを受け取る.

    Args:
        argv: ssh の引数リスト．
        input_bytes: 標準入力へ渡すバイト列（None なら何も渡さない）．
        timeout: 応答待ちの上限（秒）．超えたら ssh を止めて `_TIMEOUT_RC` を返す．

    Returns:
        実行結果．
    """
    stdin = subprocess.PIPE if input_bytes is not None else subprocess.DEVNULL
    proc = _popen(argv, stdin=stdin, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        out, err = proc.communicate(input=input_bytes, timeout=timeout)
    except subprocess.TimeoutExpired:
        proc.kill()
        out, err = proc.communicate()
        return _SshResult(_TIMEOUT_RC, out or b"", err or b"")
    return _SshResult(proc.returncode, out or b"", err or b"")


def _map_servers[T](servers: Sequence[ServerConfig], func: Callable[[ServerConfig], T]) -> list[T]:
    """サーバごとの処理を並列に実行し，servers の順で結果を返す.

    接続できないサーバの接続タイムアウトを待つ間，他のサーバの処理を止めないため並列にする．
    """
    if not servers:
        return []
    with ThreadPoolExecutor(max_workers=len(servers)) as executor:
        return list(executor.map(func, servers))


def _query_servers(
    ctx: _Context, servers: Sequence[ServerConfig], build_command: Callable[[ServerConfig], str]
) -> list[_SshResult]:
    """短い問い合わせを各サーバで並列に実行し，servers の順で結果を返す.

    コマンドは ssh を起動する前にすべて組み立てる（引数の検証エラーで一部のサーバにだけ
    問い合わせた状態にしないため）．
    """
    commands = {server.name: build_command(server) for server in servers}
    return _map_servers(
        servers,
        lambda s: _run_ssh(_ssh_argv(ctx, s, commands[s.name]), timeout=ctx.query_timeout_sec),
    )


def _run_git(project_root: Path, *git_args: str) -> str | None:
    """手元で git を実行して標準出力を返す（git が無い・失敗したら None）."""
    try:
        completed = subprocess.run(  # noqa: S603
            ["git", *git_args],
            cwd=project_root,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=_GIT_TIMEOUT_SEC,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    if completed.returncode != 0:
        return None
    return completed.stdout


def _local_git_commit(project_root: Path) -> str | None:
    """ジョブに `TIAS_GIT_COMMIT` として渡す手元の git の短縮ハッシュを求める.

    未コミットの変更（未追跡のファイルを含む）があれば末尾に `-dirty` を付ける．git が
    使えない・リポジトリでない場合は None（渡さない．ジョブ側は `nogit` とする）．
    """
    head = _run_git(project_root, "rev-parse", "--short", "HEAD")
    if head is None:
        logger.warning("git のコミットハッシュを取得できなかった（実験 ID は nogit になる）")
        return None
    status = _run_git(project_root, "status", "--porcelain")
    if status is None:
        logger.warning("git の変更状況を取得できなかった（実験 ID は nogit になる）")
        return None
    return format_git_commit(head, dirty=bool(status.strip()))


def _as_exit_code(rc: int) -> int:
    """ssh の結果の終了コードを CLI の終了コードにする（応答待ちの打ち切り等の負値は 1）."""
    return rc if rc >= 0 else 1


def _ssh_hint(ctx: _Context, server: ServerConfig, result: _SshResult) -> str:
    """よくある ssh の失敗に対処の案内を返す（該当しなければ空文字）."""
    if result.rc != _SSH_FAILURE_RC or _HOST_KEY_FAILURE not in result.stderr:
        return ""
    # BatchMode ではホスト鍵の確認に答えられないため，初回だけ手で登録してもらう
    jump = f"-J {ctx.jump_host} " if ctx.jump_host else ""
    return (
        f" → ホスト鍵が未登録: 手元で一度 ssh {jump}{server.user}@{server.address} exit を"
        "実行して登録する"
    )


def _probe_server(ctx: _Context, server: ServerConfig, threshold: GpuFreeThreshold) -> ServerProbe:
    """サーバ 1 台の tmux セッションと GPU の状況を取得する（接続できなければその旨を返す）.

    GPU の観測はサーバ側で `threshold.observe_sec` 秒かかるため，応答待ちの上限をそのぶん延ばす．
    """
    command = build_probe_command(
        server.remote_dir,
        observe_sec=threshold.observe_sec,
        sample_interval_ms=threshold.sample_interval_ms,
    )
    argv = _ssh_argv(ctx, server, command)
    result = _run_ssh(argv, timeout=ctx.query_timeout_sec + threshold.observe_sec)
    if result.rc != 0:
        return ServerProbe(server, None, _describe_failure(result) + _ssh_hint(ctx, server, result))
    try:
        output = parse_probe_output(result.stdout.decode("utf-8", errors="replace"), threshold)
    except ValueError as exc:
        return ServerProbe(server, None, str(exc))
    return ServerProbe(server, output)


def _probe_servers(
    ctx: _Context, servers: Sequence[ServerConfig], threshold: GpuFreeThreshold | None = None
) -> list[ServerProbe]:
    """複数サーバの状況を並列に取得する（servers の順で返す）.

    Args:
        ctx: 実行環境．
        servers: 対象のサーバ．
        threshold: 判定と観測の設定．None なら観測しない（ジョブの有無だけが要る push・kill
            用．瞬間値で判定する）．
    """
    if threshold is None:
        threshold = ctx.config.gpu_free.with_observe_sec(0)
    elif threshold.observe_sec > 0:
        logger.info(
            "GPU を %d 秒観測している（%d ミリ秒間隔．全サーバ並行）",
            threshold.observe_sec,
            threshold.sample_interval_ms,
        )
    return _map_servers(servers, lambda server: _probe_server(ctx, server, threshold))


def _threshold_for(ctx: _Context, observe_sec: int | None) -> GpuFreeThreshold:
    """`--observe` があれば観測秒数を差し替えた判定の設定を返す（範囲外は ValueError）."""
    threshold = ctx.config.gpu_free
    if observe_sec is None:
        return threshold
    try:
        return threshold.with_observe_sec(observe_sec)
    except ValidationError as exc:
        # 設定ファイルではなくコマンドライン引数の誤りとして伝える
        raise ValueError(f"--observe は 0〜120 の整数で指定する（指定 {observe_sec}）") from exc


def _target_servers(ctx: _Context, server_name: str | None) -> list[ServerConfig]:
    """`--server` の指定があればそのサーバ，無ければ全サーバ（優先順）を返す."""
    if server_name is not None:
        return [ctx.config.find_server(server_name)]
    return ctx.config.ordered_servers()


def _warn_unreachable(probes: Sequence[ServerProbe]) -> None:
    """接続できなかったサーバを警告し，すべて接続できなければ ValueError にする."""
    reachable = [probe for probe in probes if probe.reachable]
    if not reachable:
        details = "; ".join(f"{p.server.name}: {p.error}" for p in probes)
        raise ValueError(f"どのサーバにも接続できない（{details}）")
    for probe in probes:
        if not probe.reachable:
            logger.warning(
                "%s に接続できない（%s）．他のサーバで続行する", probe.server.name, probe.error
            )


def _report_remote_failure(result: _SshResult, server: ServerConfig) -> None:
    """ssh・リモートスクリプトに共通の失敗理由をログに出す."""
    if result.rc == EXIT_REMOTE_DIR_MISSING:
        logger.error("%s に %s が無い．先に tias-remote push する", server.name, server.remote_dir)
    elif result.rc == _SSH_FAILURE_RC:
        logger.error(
            "%s へ ssh 接続できなかった（%s）", server.name, _last_stderr_line(result.stderr)
        )
    else:
        logger.error(
            "%s でリモートコマンドが失敗した（%s）", server.name, _describe_failure(result)
        )


def _pick_single(found: Sequence[ServerConfig], *, what: str, unreachable: int) -> ServerConfig:
    """探索で見つかったサーバが 1 台であることを確かめて返す."""
    if len(found) > 1:
        names = ", ".join(server.name for server in found)
        raise ValueError(f"{what}が複数のサーバ（{names}）にある．--server で指定する")
    if not found:
        note = f"（接続できないサーバが {unreachable} 台ある）" if unreachable else ""
        raise ValueError(f"{what}がどのサーバにも見つからない{note}")
    return found[0]


# ---------------------------------------------------------------------------
# gpus / status
# ---------------------------------------------------------------------------


def _cmd_gpus(args: argparse.Namespace, ctx: _Context) -> int:
    """gpus サブコマンド: 全サーバの GPU 一覧と run が選ぶ GPU を表示する."""
    threshold = _threshold_for(ctx, args.observe)
    probes = _probe_servers(ctx, ctx.config.ordered_servers(), threshold)
    _write_stdout(format_server_report(probes, show_jobs=False))
    if not any(probe.reachable for probe in probes):
        raise ValueError("どのサーバにも接続できない")
    selection = select_gpu(probes, include_not_ready=args.no_data)
    _write_stdout(format_selection(selection) + "\n")
    if selection is None and not args.no_data:
        alternative = select_gpu(probes, include_not_ready=True)
        if alternative is not None:
            _write_stdout(
                f"  （データ未配置のサーバも含めると {alternative.server.name} の GPU "
                f"{alternative.gpu_index}．データを使わないジョブなら --no-data で選べる）\n"
            )
    return 0


def _cmd_status(args: argparse.Namespace, ctx: _Context) -> int:
    """status サブコマンド: 全サーバの実行中ジョブと GPU の状況を表示する."""
    probes = _probe_servers(ctx, ctx.config.ordered_servers(), ctx.config.gpu_free)
    _write_stdout(format_server_report(probes, show_jobs=True))
    if not any(probe.reachable for probe in probes):
        raise ValueError("どのサーバにも接続できない")
    return 0


# ---------------------------------------------------------------------------
# push
# ---------------------------------------------------------------------------


def _stream_push(argv: list[str], entries: Sequence[PushEntry]) -> int:
    """tar ストリームを ssh の標準入力へ流し，ssh の終了コードを返す."""
    proc = _popen(argv, stdin=subprocess.PIPE)
    stdin = proc.stdin
    if stdin is None:
        raise RuntimeError("ssh の標準入力を開けなかった")
    try:
        write_tar_stream(entries, stdin)
    except OSError as exc:
        # ssh が先に終了してパイプが閉じている．後始末の close 失敗は原因ではないので無視する
        with contextlib.suppress(OSError):
            stdin.close()
        rc = proc.wait()
        logger.error("送信が中断された（ssh 終了コード %d）: %s", rc, exc)
        return rc or 1
    stdin.close()
    return proc.wait()


def _push_to(ctx: _Context, probe: ServerProbe, entries: Sequence[PushEntry]) -> int:
    """サーバ 1 台へコードを送り，終了コードを返す."""
    server = probe.server
    if probe.output is not None and probe.output.jobs:
        logger.warning(
            "%s では実行中のジョブがある（%s）．上書き展開のため，実行中のジョブが"
            "新旧のコードを混ぜて読むおそれがある",
            server.name,
            ", ".join(probe.output.jobs),
        )
    total = sum(entry.size for entry in entries)
    logger.info(
        "送信: %d ファイル / %s → %s:%s",
        len(entries),
        _format_size(total),
        server.name,
        server.remote_dir,
    )
    start = time.monotonic()
    argv = _ssh_argv(ctx, server, build_push_command(server.remote_dir), stdin=True)
    rc = _stream_push(argv, entries)
    if rc != 0:
        logger.error("%s への送信に失敗した（終了コード %d）", server.name, rc)
        return rc
    logger.info("送信完了: %s（%.1f 秒）", server.name, time.monotonic() - start)
    return 0


def _collect_entries(ctx: _Context) -> list[PushEntry]:
    """送信対象を集める（空ならエラー）."""
    entries = collect_push_entries(ctx.project_root)
    if not entries:
        raise ValueError("送信するファイルが無い")
    return entries


def _cmd_push(args: argparse.Namespace, ctx: _Context) -> int:
    """push サブコマンド: コードをサーバへ送る（省略時は接続できる全サーバ）."""
    entries = _collect_entries(ctx)
    servers = _target_servers(ctx, args.server)
    if args.dry_run:
        for entry in entries:
            _write_stdout(f"  {_format_size(entry.size):>10}  {entry.arcname}\n")
        names = ", ".join(server.name for server in servers)
        _write_stdout(
            f"送信対象: {len(entries)} ファイル / "
            f"{_format_size(sum(e.size for e in entries))}（送信先の候補: {names}）\n"
            "（--dry-run のため送信していない）\n"
        )
        return 0

    probes = _probe_servers(ctx, servers)
    _warn_unreachable(probes)
    rc_total = 0
    for probe in probes:
        if probe.reachable:
            rc = _push_to(ctx, probe, entries)
            rc_total = rc_total or rc
    return rc_total


# ---------------------------------------------------------------------------
# run
# ---------------------------------------------------------------------------


def _select_gpu_target(
    args: argparse.Namespace, probes: Sequence[ServerProbe]
) -> tuple[ServerConfig, int, str] | None:
    """GPU ジョブの実行先（サーバ・GPU・GPU の状態）を選び，警告を出す.

    選べなければ状況を表示して None を返す．
    """
    selection = select_gpu(
        probes,
        server_name=args.server,
        gpu_index=args.gpu,
        force=args.force,
        include_not_ready=args.no_data,
    )
    if selection is None:
        _write_stdout(format_server_report(probes, show_jobs=True))
        scope = "全サーバ" if args.no_data else "データ配置済みのサーバ"
        logger.error(
            "%s（%s）に未使用・空き（利用者あり）の GPU が無いため，ジョブを起動しなかった",
            args.server or scope,
            "--no-data 指定" if args.no_data else "data_ready: true",
        )
        return None
    if selection.forced:
        logger.warning(
            "%s の GPU %d は%sだが，--force により使う（他のユーザーの計算と競合しうる）",
            selection.server.name,
            selection.gpu_index,
            selection.state,
        )
    shared_warning = format_shared_warning(selection)
    if shared_warning is not None:
        logger.warning("%s", shared_warning)
    return selection.server, selection.gpu_index, selection.state.value


def _select_cpu_target(
    args: argparse.Namespace, probes: Sequence[ServerProbe]
) -> tuple[ServerConfig, None, str] | None:
    """CPU ジョブ（--cpu）の実行先サーバを選ぶ（選べなければエラーを出して None）."""
    server = select_cpu_server(probes, server_name=args.server, include_not_ready=args.no_data)
    if server is None:
        logger.error(
            "接続できる%sが無いため，ジョブを起動しなかった",
            "サーバ" if args.no_data else "データ配置済み（data_ready: true）のサーバ",
        )
        return None
    return server, None, ""


def _describe_target(server: ServerConfig, gpu_index: int | None, state: str = "") -> str:
    """実行先の表示（例: `Newton の GPU 1（未使用）`・`Newton（GPU なし・CPU のみ）`）."""
    if gpu_index is None:
        return f"{server.name}（GPU なし・CPU のみ）"
    return f"{server.name} の GPU {gpu_index}" + (f"（{state}）" if state else "")


def _cmd_run(args: argparse.Namespace, ctx: _Context) -> int:
    """run サブコマンド: 空き GPU（--cpu なら GPU なし）を選び，コードを送って tmux で起動する."""
    command = list(args.command)
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        raise ValueError(
            "実行するコマンドを -- の後に指定する（例: tias-remote run --name x -- python ...）"
        )
    name = validate_job_name(args.name)
    if args.cpu:
        check_cpu_options(gpu_index=args.gpu, force=args.force, observe_sec=args.observe)
        # GPU を選ばないため観測しない（同名ジョブの確認と接続可否だけを見る）
        threshold = None
    else:
        check_selection_options(server_name=args.server, gpu_index=args.gpu, force=args.force)
        threshold = _threshold_for(ctx, args.observe)
    if args.server is not None:
        ctx.config.find_server(args.server)

    probes = _probe_servers(ctx, ctx.config.ordered_servers(), threshold)
    _warn_unreachable(probes)
    running = servers_with_session(probes, name, jobs_only=False)
    if running:
        names = ", ".join(probe.server.name for probe in running)
        raise ValueError(
            f"ジョブ名 {name!r} の tmux セッションが {names} で実行中である"
            "（--name で別名を付けるか，終了を待つ）"
        )
    if any(not probe.reachable for probe in probes):
        logger.warning("接続できないサーバでは同名ジョブの有無を確認していない")

    target = _select_cpu_target(args, probes) if args.cpu else _select_gpu_target(args, probes)
    if target is None:
        return 1
    server, gpu_index, state = target

    env = build_job_env(gpu_index, args.env, git_commit=_local_git_commit(ctx.project_root))
    script = build_job_script(
        server.remote_dir,
        name,
        command,
        uv_path=server.uv,
        server_name=server.name,
        gpu_index=gpu_index,
        env=env,
        path_env={DATA_ROOT_ENV_VAR: server.data_dir},
    )
    remote_command = build_run_command(server.remote_dir, name, script)
    if args.dry_run:
        _write_stdout(
            f"選んだ実行先: {_describe_target(server, gpu_index, state)}\n"
            f"実行するリモートコマンド:\n{remote_command}\n"
            "（--dry-run のため送信・起動していない）\n"
        )
        return 0

    if not args.no_push:
        probe = next(p for p in probes if p.server.name == server.name)
        rc = _push_to(ctx, probe, _collect_entries(ctx))
        if rc != 0:
            return rc

    result = _run_ssh(_ssh_argv(ctx, server, remote_command), timeout=ctx.query_timeout_sec)
    if result.rc == EXIT_SESSION_EXISTS:
        logger.error(
            "%s で同名のジョブ（tmux セッション）%r が実行中である（--name で別名を付ける）",
            server.name,
            name,
        )
        return result.rc
    if result.rc != 0:
        _report_remote_failure(result, server)
        return _as_exit_code(result.rc)

    log_path = f"{server.remote_dir}/{job_log_path(name)}"
    _write_stdout(
        f"ジョブを起動した: {name}（{_describe_target(server, gpu_index)}）\n"
        f"  環境変数: {' '.join(f'{k}={v}' for k, v in env.items())} "
        f"{DATA_ROOT_ENV_VAR}={server.data_dir}\n"
        f"  ログ    : {server.name}:{log_path}\n"
        "  状態確認: uv run tias-remote status\n"
        f"  ログ追従: uv run tias-remote logs {name} -f\n"
    )
    return 0


# ---------------------------------------------------------------------------
# logs / kill
# ---------------------------------------------------------------------------


def _locate_log(ctx: _Context, name: str, server_name: str | None) -> ServerConfig:
    """ジョブログがあるサーバを探す（複数・0 台ならエラー）."""
    servers = _target_servers(ctx, server_name)
    results = _query_servers(ctx, servers, lambda s: build_log_exists_command(s.remote_dir, name))
    found: list[ServerConfig] = []
    unreachable = 0
    for server, result in zip(servers, results, strict=True):
        if result.rc == 0:
            found.append(server)
        elif result.rc != EXIT_NOT_FOUND:
            unreachable += 1
            logger.warning("%s に接続できない（%s）", server.name, _describe_failure(result))
    return _pick_single(found, what=f"ジョブ {name!r} のログ", unreachable=unreachable)


def _follow(argv: list[str]) -> int:
    """ssh の出力を逐次表示し続ける（Ctrl+C で終了）.

    進捗バーは改行せず \\r で上書きすることがあるため，行単位ではなく受け取った分から表示する．
    ssh の標準入力はパイプにして開いたままにし，終了時に閉じてリモートの tail を止める．
    """
    proc = _popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    if proc.stdin is None or proc.stdout is None:
        raise RuntimeError("ssh の標準入出力を開けなかった")
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    interrupted = False
    try:
        while chunk := proc.stdout.read1(65536):
            _write_stdout(decoder.decode(chunk))
    except KeyboardInterrupt:
        interrupted = True
        logger.info("（追従を終了した．ジョブは継続している）")
    finally:
        # stdin を閉じるとリモートの cat が EOF を受けて tail を止める．既に ssh が
        # 終了していれば close は失敗しうるが，後始末なので無視する
        with contextlib.suppress(OSError):
            proc.stdin.close()
    rc = proc.wait()
    # Ctrl+C は追従の正常な終わらせ方（Windows では ssh も同時に止まり非 0 を返す）
    return 0 if interrupted else rc


def _cmd_logs(args: argparse.Namespace, ctx: _Context) -> int:
    """logs サブコマンド: ジョブのログを表示する（ジョブのあるサーバを自動で探す）."""
    name = validate_job_name(args.name)
    if args.lines < 0:
        raise ValueError("--lines は 0 以上を指定する")
    server = _locate_log(ctx, name, args.server)
    command = build_logs_command(server.remote_dir, name, lines=args.lines, follow=args.follow)
    if args.follow:
        logger.info("%s のジョブ %s のログを追従する（Ctrl+C で終了）", server.name, name)
        return _follow(_ssh_argv(ctx, server, command, stdin=True))
    result = _run_ssh(_ssh_argv(ctx, server, command))
    _write_stdout(result.stdout.decode("utf-8", errors="replace"))
    if result.rc == EXIT_NOT_FOUND:
        logger.error("%s にジョブ %r のログが無い", server.name, name)
    elif result.rc != 0:
        _report_remote_failure(result, server)
    return _as_exit_code(result.rc)


def _cmd_kill(args: argparse.Namespace, ctx: _Context) -> int:
    """kill サブコマンド: ジョブを止める（ジョブのあるサーバを自動で探す）."""
    name = validate_job_name(args.name)
    probes = _probe_servers(ctx, _target_servers(ctx, args.server))
    _warn_unreachable(probes)
    found = servers_with_session(probes, name, jobs_only=True)
    if not found and servers_with_session(probes, name, jobs_only=False):
        raise ValueError(
            f"{name!r} は tias-remote のジョブではない tmux セッションである（止めない）"
        )
    unreachable = sum(1 for probe in probes if not probe.reachable)
    server = _pick_single(
        [probe.server for probe in found], what=f"実行中のジョブ {name!r}", unreachable=unreachable
    )
    result = _run_ssh(
        _ssh_argv(ctx, server, build_kill_command(server.remote_dir, name)),
        timeout=ctx.query_timeout_sec,
    )
    if result.rc == EXIT_NO_SESSION:
        logger.error("%s のジョブ %r は実行中ではない（既に終了した可能性）", server.name, name)
        return result.rc
    if result.rc != 0:
        _report_remote_failure(result, server)
        return _as_exit_code(result.rc)
    _write_stdout(f"ジョブ {name} を終了した（{server.name}）\n")
    return 0


# ---------------------------------------------------------------------------
# pull
# ---------------------------------------------------------------------------


def _locate_run(ctx: _Context, run_id: str, server_name: str | None) -> tuple[ServerConfig, bytes]:
    """RUN_ID のディレクトリがあるサーバを探し，そのファイル一覧を返す（複数・0 台ならエラー）."""
    servers = _target_servers(ctx, server_name)
    results = _query_servers(ctx, servers, lambda s: build_pull_list_command(s.remote_dir, run_id))
    found: list[tuple[ServerConfig, bytes]] = []
    unreachable = 0
    for server, result in zip(servers, results, strict=True):
        if result.rc == 0:
            found.append((server, result.stdout))
        elif result.rc == EXIT_UNSAFE_PATH:
            raise ValueError(
                f"{server.name} の {REMOTE_RUNS_DIR}/{run_id} がリンクであるため回収しない"
            )
        elif result.rc != EXIT_NOT_FOUND:
            # 一覧の失敗理由（find の標準エラー）はファイル名を含みうるので表示しない
            unreachable += 1
            reason = (
                "応答がタイムアウトした" if result.rc == _TIMEOUT_RC else f"終了コード {result.rc}"
            )
            logger.warning("%s で実験結果を探せなかった（%s）", server.name, reason)
    server = _pick_single([s for s, _ in found], what=f"実験 {run_id!r}", unreachable=unreachable)
    listing = next(out for s, out in found if s.name == server.name)
    return server, listing


def _receive_tar(argv: list[str], payload: bytes, destination: BinaryIO) -> _SshResult:
    """ssh の標準入力へ payload を送り，標準出力（tar）を destination へ書き写す.

    書き込みと読み出しを別スレッドにする（送る一覧が大きいとき，リモートの tar の出力で
    パイプが詰まって互いに待つのを避けるため）．標準エラーは表示しない（リモートの tar の
    メッセージはファイル名を含みうるため）．
    """
    proc = _popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if proc.stdin is None or proc.stdout is None or proc.stderr is None:
        raise RuntimeError("ssh の標準入出力を開けなかった")
    stdin, stderr_pipe = proc.stdin, proc.stderr
    stderr_chunks: list[bytes] = []

    def _write_payload() -> None:
        # ssh が先に終了していれば書き込みは失敗する．終了コードで失敗を報告するので無視する
        with contextlib.suppress(OSError):
            stdin.write(payload)
        with contextlib.suppress(OSError):
            stdin.close()

    writer = threading.Thread(target=_write_payload, daemon=True)
    reader = threading.Thread(target=lambda: stderr_chunks.append(stderr_pipe.read()), daemon=True)
    writer.start()
    reader.start()
    copy_stream(proc.stdout, destination)
    writer.join()
    reader.join()
    rc = proc.wait()
    return _SshResult(rc, b"", b"".join(stderr_chunks))


def _cmd_pull(args: argparse.Namespace, ctx: _Context) -> int:
    """pull サブコマンド: 実験結果を回収ホワイトリストに従って持ち帰る."""
    run_id = validate_run_id(args.run_id)
    server, listing = _locate_run(ctx, run_id, args.server)
    plan = plan_pull(parse_find_listing(listing), run_id, with_checkpoints=args.with_checkpoints)
    if not plan.files:
        raise ValueError(
            f"{server.name} の実験 {run_id!r} に回収対象のファイルが無い"
            f"（回収しなかったファイル {plan.skipped} 件）"
        )

    dest = ctx.project_root / REMOTE_RUNS_DIR
    logger.info("取得: %s の %s/%s → %s", server.name, REMOTE_RUNS_DIR, run_id, dest / run_id)
    payload = b"".join(path.encode("utf-8", "surrogateescape") + b"\0" for path in plan.files)
    argv = _ssh_argv(ctx, server, build_pull_tar_command(server.remote_dir), stdin=True)
    start = time.monotonic()
    with tempfile.TemporaryFile() as archive:
        result = _receive_tar(argv, payload, archive)
        if result.rc == _TAR_WARNING_RC:
            logger.warning("一部のファイルが読み取り中に変化した（ジョブが実行中の可能性）")
        elif result.rc == _SSH_FAILURE_RC:
            logger.error(
                "%s へ ssh 接続できなかった（%s）", server.name, _last_stderr_line(result.stderr)
            )
            return result.rc
        elif result.rc != 0:
            logger.error("%s で結果の書き出しに失敗した（終了コード %d）", server.name, result.rc)
            return result.rc
        try:
            extracted = extract_pull_archive(
                archive, dest, run_id, with_checkpoints=args.with_checkpoints
            )
        except tarfile.TarError as exc:
            raise ValueError("受け取ったアーカイブが壊れているため，何も展開せず中止した") from exc

    if extracted.files != len(plan.files):
        logger.warning(
            "回収予定 %d 件のうち %d 件しか届かなかった", len(plan.files), extracted.files
        )
    skipped = f"回収しなかったファイル: {plan.skipped} 件"
    if plan.skipped_checkpoints:
        skipped += (
            f"（うち checkpoints/*.pt {plan.skipped_checkpoints} 件は"
            " --with-checkpoints で回収できる）"
        )
    _write_stdout(
        f"取得完了: {extracted.files} ファイル / {_format_size(extracted.bytes)}"
        f"（{time.monotonic() - start:.1f} 秒）→ {REMOTE_RUNS_DIR}/{run_id}/\n{skipped}\n"
    )
    return 0


# ---------------------------------------------------------------------------
# 引数
# ---------------------------------------------------------------------------


def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace:
    """コマンドライン引数をパースして返す."""
    parser = argparse.ArgumentParser(
        prog="tias-remote",
        description=(
            "研究室の GPU サーバ（configs/remote.yaml）で空いている GPU を選んで実行する"
            "（コード送信・ジョブ起動・状態確認・結果回収）．"
        ),
        epilog=_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument(
        "--config",
        type=Path,
        default=None,
        metavar="PATH",
        help=f"サーバ設定 YAML（既定: プロジェクトルートの {DEFAULT_CONFIG_PATH.as_posix()}）．",
    )
    server_opt = argparse.ArgumentParser(add_help=False)
    server_opt.add_argument(
        "--server", default=None, metavar="NAME", help="対象のサーバ（既定: 全サーバから探す）．"
    )
    sub = parser.add_subparsers(dest="subcommand", required=True, metavar="SUBCOMMAND")

    p_gpus = sub.add_parser(
        "gpus",
        parents=[common],
        help="全サーバの GPU 一覧（状態・使い方の型・利用者）と run が選ぶ GPU を表示する",
        description=(
            "各 GPU を一定時間観測し，観測中の最大値で 未使用（閾値内・計算プロセス無し）／"
            "空き（利用者あり）／使用中／状態不明 に分けて表示する．run は全候補サーバを通して"
            "まず未使用を，無ければ空き（利用者あり）を選ぶ．"
        ),
    )
    p_gpus.add_argument(
        "--observe",
        type=int,
        default=None,
        metavar="SEC",
        help="観測秒数（設定の observe_sec を上書き．0 で瞬間値のみ．0〜120）．",
    )
    p_gpus.add_argument(
        "--no-data",
        action="store_true",
        help="データ未配置（data_ready: false）のサーバも run の候補として選ぶ．",
    )

    sub.add_parser("status", parents=[common], help="全サーバの実行中ジョブと GPU の状況を表示する")

    p_push = sub.add_parser(
        "push",
        parents=[common, server_opt],
        help="コードをサーバへ送る（省略時は接続できる全サーバ）",
        description=(
            "リポジトリを tar over ssh でサーバの remote_dir へ上書き展開する．"
            ".git / .venv / キャッシュ / data / runs / .claude は送らない．"
        ),
    )
    p_push.add_argument(
        "--dry-run", action="store_true", help="送信せず，対象ファイルの一覧だけ表示する．"
    )

    p_run = sub.add_parser(
        "run",
        parents=[common, server_opt],
        help="GPU を選び（未使用を最優先），コードを送って tmux でジョブを起動する",
        description=(
            "GPU を選び（data_ready のサーバ全体を通して，まず未使用を，無ければ"
            "空き（利用者あり）を priority 順・GPU 番号順に探す），"
            "そのサーバへコードを送ってから remote_dir で `uv run <COMMAND>` を tmux で起動する．"
            f"出力は remote_dir/{REMOTE_LOG_DIR}/<NAME>.log に追記され，接続が切れても継続する．"
            "コマンドは -- の後に書く．"
        ),
    )
    p_run.add_argument(
        "--name",
        required=True,
        metavar="NAME",
        help="ジョブ名（tmux セッション名・ログ名）．英数字で始まり英数字・_・-，64 文字以内．",
    )
    p_run.add_argument(
        "--gpu",
        type=int,
        default=None,
        metavar="I",
        help="使う GPU 番号（--server が必須．使用中・状態不明ならエラー）．",
    )
    p_run.add_argument(
        "--force",
        action="store_true",
        help="--gpu で指定した GPU が使用中・状態不明でも使う．",
    )
    p_run.add_argument(
        "--cpu",
        action="store_true",
        help=(
            "GPU を使わないジョブ（manifest の生成・データセットの集計等）として起動する．"
            "GPU を観測・選択せず，CUDA_VISIBLE_DEVICES を空にする．サーバは --server，"
            "省略時は候補サーバを priority 順に見て最初に接続できたもの．"
            "--gpu・--force・--observe とは同時に指定できない．"
        ),
    )
    p_run.add_argument(
        "--no-data",
        action="store_true",
        help="データ未配置（data_ready: false）のサーバも候補にする．",
    )
    p_run.add_argument("--no-push", action="store_true", help="起動前にコードを送らない．")
    p_run.add_argument(
        "--observe",
        type=int,
        default=None,
        metavar="SEC",
        help="GPU を選ぶ前の観測秒数（設定の observe_sec を上書き．0 で瞬間値のみ）．",
    )
    p_run.add_argument(
        "--env",
        action="append",
        default=[],
        metavar="KEY=VALUE",
        help=(
            "ジョブに追加で渡す環境変数．複数指定可（tias-remote が設定する CUDA_VISIBLE_DEVICES・"
            "TIAS_DATA_ROOT・TIAS_GIT_COMMIT 等は指定できない）．"
        ),
    )
    p_run.add_argument(
        "--dry-run",
        action="store_true",
        help="GPU を選ぶところまで行い，送信・起動せずにリモートコマンドを表示する．",
    )
    p_run.add_argument(
        "command",
        nargs=argparse.REMAINDER,
        metavar="-- COMMAND",
        help="サーバで実行するコマンド（uv run の後ろに付く部分）．",
    )

    p_logs = sub.add_parser(
        "logs", parents=[common, server_opt], help="ジョブのログを表示する（サーバを自動で探す）"
    )
    p_logs.add_argument("name", metavar="NAME", help="ジョブ名．")
    p_logs.add_argument(
        "--lines", "-n", type=int, default=40, metavar="N", help="表示する末尾の行数（既定 40）．"
    )
    p_logs.add_argument(
        "--follow", "-f", action="store_true", help="追記を追従表示する（Ctrl+C で終了）．"
    )

    p_kill = sub.add_parser(
        "kill", parents=[common, server_opt], help="ジョブを止める（サーバを自動で探す）"
    )
    p_kill.add_argument("name", metavar="NAME", help="ジョブ名．")

    p_pull = sub.add_parser(
        "pull",
        parents=[common, server_opt],
        help="実験結果を回収ホワイトリストに従って持ち帰る",
        description=(
            "サーバの runs/<RUN_ID>/ のうち回収ホワイトリスト（config.yaml・history.csv・"
            "metrics.json・confusion_matrix.csv・log.txt・figures/*.png）に一致するファイルだけを"
            "手元の runs/<RUN_ID>/ へ持ち帰る（同名ファイルは上書き）．"
        ),
    )
    p_pull.add_argument("run_id", metavar="RUN_ID", help="実験 ID（runs/ 直下のディレクトリ名）．")
    p_pull.add_argument(
        "--with-checkpoints", action="store_true", help="checkpoints/*.pt も持ち帰る．"
    )

    return parser.parse_args(argv)


_HANDLERS: dict[str, Callable[[argparse.Namespace, _Context], int]] = {
    "gpus": _cmd_gpus,
    "status": _cmd_status,
    "push": _cmd_push,
    "run": _cmd_run,
    "logs": _cmd_logs,
    "kill": _cmd_kill,
    "pull": _cmd_pull,
}


def _build_context(args: argparse.Namespace) -> _Context:
    """設定を読み込み，実行環境を組み立てる."""
    project_root = find_project_root(Path.cwd())
    config_path = args.config if args.config is not None else project_root / DEFAULT_CONFIG_PATH
    config = load_remote_config(config_path)
    return _Context(
        config=config,
        ssh=resolve_ssh_executable(os.environ),
        jump_host=resolve_jump_host(config, os.environ),
        project_root=project_root,
    )


def main(argv: Sequence[str] | None = None) -> int:
    """tias-remote コマンドのメインエントリポイント.

    Args:
        argv: 引数リスト．None のとき sys.argv[1:] を使用．

    Returns:
        終了コード（成功時 0．ssh・リモートの失敗時はその終了コード，
        入力エラー・空き GPU 無しは 1）．
    """
    args = _parse_args(argv)
    _configure_logging()
    try:
        ctx = _build_context(args)
        if not is_self_check_disabled(os.environ):
            self_server = detect_self_server(ctx.config, socket.gethostname())
            if self_server is not None:
                logger.error(
                    "ここはサーバ %s である（ホスト名が一致）．"
                    "tias-remote は手元の PC から実行する．"
                    "サーバ上で直接回すなら tmux の中で uv run ... を実行する．"
                    "判定が誤りなら環境変数 %s=1 で無効にできる",
                    self_server.name,
                    ALLOW_SELF_ENV_VAR,
                )
                return 1
        if ctx.jump_host is None and JUMP_ENV_VAR in os.environ:
            logger.info("%s が空のため，踏み台を使わず直接接続する", JUMP_ENV_VAR)
        return _HANDLERS[args.subcommand](args, ctx)
    except ValidationError as exc:
        logger.error("サーバ設定が不正である:\n%s", exc)
        return 1
    except ValueError as exc:
        logger.error("%s", exc)
        return 1
    except KeyboardInterrupt:
        logger.error("中断した")
        return 130


if __name__ == "__main__":
    # uv の生成するランチャ（tias-remote.exe）が Windows の Smart App Control に
    # ブロックされた場合の代替として `uv run python -m tias_judge_model.cli.remote` を使える
    sys.exit(main())