"""研究室サーバ上で実行するシェルコマンド文字列を組み立てる純関数群．

`ssh <host> '<command>'` で送る文字列を作る．リモートのログインシェル（Ubuntu では bash）が
解釈する前提で POSIX sh の範囲で書き，ジョブ本体だけは `${PIPESTATUS[0]}` を使うため
`bash -c` で包む．ユーザー由来の値（パス・コマンド引数）はすべて `shlex.quote` でクォートし，
ジョブ名・GPU 指定・環境変数名は書式を検証してから埋め込む．

Windows から ssh へ渡す引数に改行を含めないため，複数の文は `; ` で 1 行につなぐ．
リモート側が出すメッセージは ASCII に限る（ssh の stderr は Windows コンソールへ素通しされ，
コードページによっては日本語が化けるため．日本語の説明は呼び出し側の CLI が出す）．
"""

from __future__ import annotations

import re
import shlex
from collections.abc import Mapping, Sequence
from pathlib import PurePosixPath

REMOTE_OUTPUT_DIR = "output"
REMOTE_LOG_DIR = "output/_remote_logs"

# リモートスクリプトが返す終了コード．ssh 自身の 255・tar の 1/2 と衝突しない値にする
EXIT_REMOTE_DIR_MISSING = 90
EXIT_SESSION_EXISTS = 91
EXIT_NOT_FOUND = 92
EXIT_NO_SESSION = 93

# tmux のセッション名に使えない "." ":" を避け，ログのファイル名にもそのまま使える文字に限る
_JOB_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$")
_GPU_SPEC_RE = re.compile(r"^\d+(,\d+)*$")
_ENV_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")

_NOW = "\"$(date '+%Y-%m-%d %H:%M:%S')\""
_STATUS_LOG_COUNT = 10


def quote_remote_path(path: str) -> str:
    """リモートのパスをシェル用にクォートする．

    先頭の `~` はリモートの `$HOME` に展開させる必要があるため，`~/...` は
    `"$HOME"/'...'` の形に変換する（丸ごとクォートするとチルダ展開が起きない）．

    Args:
        path: リモートのパス（`~/...`・絶対パス・ホーム基準の相対パス）．

    Returns:
        シェルにそのまま埋め込める文字列．

    Raises:
        ValueError: 空文字，または `~ユーザー名/...` 形式の場合．
    """
    if not path or not path.strip():
        raise ValueError("リモートパスが空です．")
    if path == "~":
        return '"$HOME"'
    if path.startswith("~/"):
        rest = path[2:]
        return f'"$HOME"/{shlex.quote(rest)}' if rest else '"$HOME"'
    if path.startswith("~"):
        msg = (
            f"~ユーザー名 形式のリモートパスには対応していません: {path!r}"
            "（~/... か絶対パスを指定）"
        )
        raise ValueError(msg)
    return shlex.quote(path)


def validate_host(host: str) -> str:
    """SSH ホスト名を検証して返す（ssh のオプションとして解釈されるのを防ぐ）．

    Args:
        host: `~/.ssh/config` の Host 名，またはホスト名・IP．

    Returns:
        検証済みのホスト名．

    Raises:
        ValueError: 空・`-` 始まり・空白を含む場合．
    """
    if not host or host.startswith("-") or any(c.isspace() for c in host):
        raise ValueError(f"SSH ホスト名が不正です: {host!r}")
    return host


def validate_job_name(name: str) -> str:
    """ジョブ名（tmux セッション名・ログのファイル名）を検証して返す．

    Args:
        name: ジョブ名．

    Returns:
        検証済みのジョブ名．

    Raises:
        ValueError: 英数字で始まり英数字・`_`・`-` のみ（64 文字以内）でない場合．
    """
    if not _JOB_NAME_RE.match(name):
        msg = f"ジョブ名が不正です: {name!r}（英数字で始まり，英数字・_・- のみ，64 文字以内）"
        raise ValueError(msg)
    return name


def validate_gpu_spec(spec: str) -> str:
    """CUDA_VISIBLE_DEVICES に渡す GPU 番号の指定を検証して返す．

    Args:
        spec: `0` や `0,2` のようなカンマ区切りの GPU 番号．

    Returns:
        検証済みの指定文字列．

    Raises:
        ValueError: 書式が不正な場合．
    """
    if not _GPU_SPEC_RE.match(spec):
        raise ValueError(f"GPU 指定が不正です: {spec!r}（例: 0 / 0,2）")
    return spec


def parse_env_assignment(text: str) -> tuple[str, str]:
    """`KEY=VALUE` 形式の環境変数指定を分解する．

    Args:
        text: `KEY=VALUE` 形式の文字列（VALUE は空でもよい）．

    Returns:
        (KEY, VALUE)．

    Raises:
        ValueError: `=` が無い，または KEY が環境変数名として不正な場合．
    """
    key, sep, value = text.partition("=")
    if not sep or not _ENV_KEY_RE.match(key):
        raise ValueError(f"環境変数の指定が不正です: {text!r}（KEY=VALUE 形式）")
    return key, value


def build_ssh_argv(ssh: str, host: str, remote_command: str, *, tty: bool = False) -> list[str]:
    """ssh を起動する引数リストを返す．

    Args:
        ssh: ssh 実行ファイル．
        host: 接続先ホスト．
        remote_command: リモートで実行するシェルコマンド（1 引数として渡す）．
        tty: True なら `-t` で疑似端末を割り当てる（tmux attach 用）．

    Returns:
        subprocess に渡す引数リスト．
    """
    argv = [ssh]
    if tty:
        argv.append("-t")
    argv += [validate_host(host), remote_command]
    return argv


def _cd_command(remote_dir: str) -> str:
    """リモートディレクトリへ移動し，無ければ専用の終了コードで抜ける文を返す．"""
    quoted = quote_remote_path(remote_dir)
    return (
        f"cd {quoted} || {{ echo 'sm-remote: remote directory not found:' {quoted} >&2; "
        f"exit {EXIT_REMOTE_DIR_MISSING}; }}"
    )


def _job_log_path(name: str) -> str:
    """ジョブログのパス（プロジェクトルート基準．クォート前）を返す．name は検証済みとする．"""
    return f"{REMOTE_LOG_DIR}/{name}.log"


def _tmux_target(name: str) -> str:
    """tmux の `-t` に渡すセッション指定をクォートして返す．name は検証済みとする．

    `=` を前置して完全一致にする（付けないと tmux は前方一致で別のジョブを選びうる）．
    """
    return shlex.quote(f"={name}")


def build_push_command(remote_dir: str, *, compress: bool) -> str:
    """標準入力の tar ストリームをリモートディレクトリへ展開するコマンドを返す．

    展開前に tmux セッションが動いていれば警告を出す（実行中のジョブが途中から
    新しいコードを読み込み，結果が混ざるおそれがあるため）．

    Args:
        remote_dir: 展開先（リモートのプロジェクトルート）．
        compress: True なら gzip 圧縮ストリームとして展開する．

    Returns:
        リモートで実行するシェルコマンド．
    """
    quoted = quote_remote_path(remote_dir)
    flags = "-xzf" if compress else "-xf"
    warn = (
        "if tmux ls >/dev/null 2>&1; then "
        "echo 'sm-remote: warning: tmux sessions are running"
        " (running jobs may pick up the new code):' >&2; tmux ls >&2; fi"
    )
    return f"{warn}; mkdir -p {quoted} && tar {flags} - -C {quoted}"


def build_list_files_command(
    remote_dir: str, dirs: Sequence[str], skip_dir_names: Sequence[str]
) -> str:
    """リモートの指定ディレクトリ配下のファイル一覧を NUL 区切りで出力するコマンドを返す．

    Args:
        remote_dir: リモートのプロジェクトルート．
        dirs: 一覧を取るディレクトリ（プロジェクトルート基準）．
        skip_dir_names: 探索しないディレクトリ名（`__pycache__` 等）．

    Returns:
        リモートで実行するシェルコマンド．
    """
    names = " ".join(shlex.quote(d) for d in dirs)
    prune = " -o ".join(f"-name {shlex.quote(n)}" for n in skip_dir_names)
    return (
        f"{_cd_command(remote_dir)}; for d in {names}; do "
        f'if [ -d "$d" ]; then find "$d" -type d \\( {prune} \\) -prune -o -type f -print0; fi; '
        "done"
    )


def build_remove_files_command(remote_dir: str) -> str:
    """標準入力の NUL 区切りパス一覧をリモートで削除するコマンドを返す．

    Args:
        remote_dir: リモートのプロジェクトルート（パスはこの基準の相対パス）．

    Returns:
        リモートで実行するシェルコマンド．
    """
    return f"{_cd_command(remote_dir)}; xargs -0 rm -f --"


def build_job_script(
    remote_dir: str,
    name: str,
    command: Sequence[str],
    *,
    uv_path: str,
    use_uv: bool,
    env: Mapping[str, str],
) -> str:
    """tmux セッション内で bash が実行するジョブ本体のスクリプトを返す．

    出力は `output/_remote_logs/<name>.log` に追記し（前回の同名ジョブのログも残す），
    開始行と終了コード行で挟む．標準入力は /dev/null にする（入力待ちで止まらないように）．
    uv のディレクトリを PATH に足すのは，スイープスクリプトが `uv run sm-render` を
    名前で呼ぶため（非ログインの ssh では ~/.local/bin が PATH に入らない）．

    Args:
        remote_dir: リモートのプロジェクトルート．
        name: ジョブ名．
        command: 実行するコマンドと引数．
        uv_path: リモートの uv 実行ファイル．
        use_uv: True なら `uv run` を前置する．
        env: ジョブに渡す環境変数．

    Returns:
        bash スクリプト（1 行）．

    Raises:
        ValueError: コマンドが空，またはジョブ名・環境変数名が不正な場合．
    """
    if not command:
        raise ValueError("実行するコマンドが空です．")
    log = _job_log_path(validate_job_name(name))
    for key in env:
        if not _ENV_KEY_RE.match(key):
            raise ValueError(f"環境変数名が不正です: {key!r}")

    uv_quoted = quote_remote_path(uv_path)
    run_argv = ([uv_quoted, "run"] if use_uv else []) + [shlex.quote(arg) for arg in command]
    display = shlex.join(
        [f"{k}={v}" for k, v in env.items()] + (["uv", "run"] if use_uv else []) + list(command)
    )
    statements = [
        _cd_command(remote_dir),
        f"log={shlex.quote(log)}",
        f"printf '=== [%s] start on %s: %s\\n' {_NOW} \"$(hostname)\" {shlex.quote(display)}"
        ' >> "$log"',
        f'export PATH="$(dirname -- {uv_quoted}):$PATH"',
        *[f"export {k}={shlex.quote(v)}" for k, v in env.items()],
        f'{" ".join(run_argv)} < /dev/null 2>&1 | tee -a "$log"',
        "rc=${PIPESTATUS[0]}",
        f'printf \'=== [%s] exit %s\\n\' {_NOW} "$rc" >> "$log"',
        'exit "$rc"',
    ]
    return "; ".join(statements)


def build_run_command(remote_dir: str, name: str, job_script: str) -> str:
    """ジョブ本体を tmux の新規セッション（デタッチ状態）で起動するコマンドを返す．

    同名セッションが既にあれば起動せず `EXIT_SESSION_EXISTS` で終わる．
    tmux のサーバはデーモン化して ssh の入出力を閉じるため，ssh はすぐ戻り，
    接続が切れてもジョブは継続する．

    Args:
        remote_dir: リモートのプロジェクトルート．
        name: ジョブ名（tmux セッション名）．
        job_script: `build_job_script` が返したスクリプト．

    Returns:
        リモートで実行するシェルコマンド．
    """
    name = validate_job_name(name)
    target = _tmux_target(name)
    inner = "bash -c " + shlex.quote(job_script)
    statements = [
        _cd_command(remote_dir),
        f"mkdir -p {REMOTE_LOG_DIR} || exit 1",
        f"if tmux has-session -t {target} 2>/dev/null; then "
        f"echo 'sm-remote: tmux session already exists:' {shlex.quote(name)} >&2; "
        f"exit {EXIT_SESSION_EXISTS}; fi",
        f'tmux new-session -d -s {shlex.quote(name)} -c "$PWD" {shlex.quote(inner)}',
    ]
    return "; ".join(statements)


def _last_line_command(path_expr: str) -> str:
    """ログの最後の 1 行を表示するパイプラインを返す．

    mitsuba の進捗バーは \\r で上書きするので，最後の \\r 以降だけを 1 行として見せる．
    """
    return f"tail -n 1 {path_expr} | tr '\\r' '\\n' | grep -v '^$' | tail -n 1 | cut -c1-100"


def build_status_command(remote_dir: str) -> str:
    """tmux セッション・最近のジョブログ・GPU・負荷を表示するコマンドを返す．

    Args:
        remote_dir: リモートのプロジェクトルート．

    Returns:
        リモートで実行するシェルコマンド．
    """
    last_line = _last_line_command('"$f"')
    log_loop = (
        f"for f in $(ls -t {REMOTE_LOG_DIR}/*.log 2>/dev/null | head -n {_STATUS_LOG_COUNT}); do "
        'n=$(basename "$f" .log); '
        'if tmux has-session -t "=$n" 2>/dev/null; then s=RUNNING; else s=ended; fi; '
        f"last=$({last_line}); "
        'printf \'%-28s %-8s %s\\n\' "$n" "$s" "$last"; '
        "done"
    )
    gpu = (
        "if command -v nvidia-smi >/dev/null 2>&1; then "
        "nvidia-smi --query-gpu=index,name,utilization.gpu,memory.used,memory.total"
        " --format=csv,noheader; "
        "else echo '(nvidia-smi not found)'; fi"
    )
    statements = [
        _cd_command(remote_dir),
        "echo '== tmux sessions (running jobs) =='",
        "tmux ls 2>/dev/null || echo '(none)'",
        "echo",
        f"echo '== recent job logs ({REMOTE_LOG_DIR}) =='",
        log_loop,
        "echo",
        "echo '== GPU (index, name, util, mem used, mem total) =='",
        gpu,
        "echo",
        "echo '== load average =='",
        "uptime",
    ]
    return "; ".join(statements)


def build_logs_command(remote_dir: str, name: str, *, lines: int, follow: bool) -> str:
    """ジョブログの末尾を表示する（follow なら追従する）コマンドを返す．

    追従モードでは，ssh の標準入力が閉じたら（手元で Ctrl+C・終了したら）tail を止める．
    疑似端末なしの ssh では切断しても tail に SIGHUP が届かず，ログが伸びない限り
    リモートに tail が残り続けるため．

    Args:
        remote_dir: リモートのプロジェクトルート．
        name: ジョブ名．
        lines: 表示する末尾の行数．
        follow: True なら追記を追従する．

    Returns:
        リモートで実行するシェルコマンド．
    """
    log = shlex.quote(_job_log_path(validate_job_name(name)))
    check = (
        f"[ -f {log} ] || {{ echo 'sm-remote: log not found:' {log} >&2; exit {EXIT_NOT_FOUND}; }}"
    )
    if follow:
        tail = f'tail -n {int(lines)} -F {log} & t=$!; cat >/dev/null; kill "$t" 2>/dev/null'
    else:
        tail = f"tail -n {int(lines)} {log}"
    return "; ".join([_cd_command(remote_dir), check, tail])


def build_attach_command(name: str) -> str:
    """ジョブの tmux セッションに接続するコマンドを返す（ssh -t で使う）．

    Args:
        name: ジョブ名．

    Returns:
        リモートで実行するシェルコマンド．
    """
    return f"tmux attach-session -t {_tmux_target(validate_job_name(name))}"


def build_attach_check_command(remote_dir: str, name: str) -> str:
    """attach の前に，ジョブの tmux セッションがまだあるかを確かめるコマンドを返す．

    セッションがあれば 0 で終わる．無ければログの最終行（`=== [...] exit N` 等）を
    標準出力に出して `EXIT_NO_SESSION`，ログも無ければ `EXIT_NOT_FOUND` で終わる．
    `ssh -t` の対話セッションは出力を受け取れないため，事前チェックは別の ssh で行う．

    Args:
        remote_dir: リモートのプロジェクトルート．
        name: ジョブ名．

    Returns:
        リモートで実行するシェルコマンド．
    """
    name = validate_job_name(name)
    log = shlex.quote(_job_log_path(name))
    statements = [
        f"if tmux has-session -t {_tmux_target(name)} 2>/dev/null; then exit 0; fi",
        _cd_command(remote_dir),
        f"if [ -f {log} ]; then {_last_line_command(log)}; exit {EXIT_NO_SESSION}; fi",
        f"exit {EXIT_NOT_FOUND}",
    ]
    return "; ".join(statements)


def build_kill_command(remote_dir: str, name: str) -> str:
    """ジョブの tmux セッションを終了し，ログにその旨を追記するコマンドを返す．

    tmux がセッションを閉じるとジョブの bash は SIGHUP で終わり，自分では終了行を
    書けないため，終了前にここで追記する．

    Args:
        remote_dir: リモートのプロジェクトルート．
        name: ジョブ名．

    Returns:
        リモートで実行するシェルコマンド．
    """
    name = validate_job_name(name)
    target = _tmux_target(name)
    log = shlex.quote(_job_log_path(name))
    statements = [
        f"tmux has-session -t {target} 2>/dev/null || "
        f"{{ echo 'sm-remote: no such tmux session:' {shlex.quote(name)} >&2; "
        f"exit {EXIT_NO_SESSION}; }}",
        f"if cd {quote_remote_path(remote_dir)} 2>/dev/null && [ -f {log} ]; then "
        f"printf '=== [%s] killed by sm-remote kill\\n' {_NOW} >> {log}; fi",
        f"tmux kill-session -t {target}",
    ]
    return "; ".join(statements)


def normalize_output_subpath(text: str) -> str:
    """pull 対象の指定を `output/<サブパス>` 形式に正規化する．

    `offaxis_led_sweep`・`output/offaxis_led_sweep`・`output\\offaxis_led_sweep`（Windows 区切り）
    をいずれも `output/offaxis_led_sweep` にする．output/ の外や output/ 全体は指定できない．

    Args:
        text: ユーザーが指定したパス．

    Returns:
        リモートのプロジェクトルート基準の相対パス（POSIX 区切り）．

    Raises:
        ValueError: 空・絶対パス・`..` を含む・output/ 全体を指す場合．
    """
    cleaned = text.strip().replace("\\", "/")
    if not cleaned:
        raise ValueError("取得するパスが空です．")
    if cleaned.startswith("/") or _WINDOWS_DRIVE_RE.match(cleaned):
        raise ValueError(f"絶対パスは指定できません: {text!r}（output/ 配下の相対パスを指定）")
    parts = [p for p in PurePosixPath(cleaned).parts if p not in ("", ".")]
    if ".." in parts:
        raise ValueError(f"'..' を含むパスは指定できません: {text!r}")
    if parts and parts[0] == REMOTE_OUTPUT_DIR:
        parts = parts[1:]
    if not parts:
        raise ValueError("output/ 全体は指定できません．サブディレクトリを指定してください．")
    return "/".join([REMOTE_OUTPUT_DIR, *parts])


def build_pull_command(remote_dir: str, rel_paths: Sequence[str], *, compress: bool) -> str:
    """指定パスを tar ストリームとして標準出力へ書き出すコマンドを返す．

    どれか 1 つでも存在しなければ tar を始める前に `EXIT_NOT_FOUND` で終わる．

    Args:
        remote_dir: リモートのプロジェクトルート．
        rel_paths: `normalize_output_subpath` で正規化済みのパス．
        compress: True なら gzip 圧縮する．

    Returns:
        リモートで実行するシェルコマンド．
    """
    if not rel_paths:
        raise ValueError("取得するパスが指定されていません．")
    quoted = " ".join(shlex.quote(p) for p in rel_paths)
    flags = "-czf" if compress else "-cf"
    statements = [
        _cd_command(remote_dir),
        f'for p in {quoted}; do [ -e "$p" ] || '
        f"{{ echo 'sm-remote: not found on remote:' \"$p\" >&2; exit {EXIT_NOT_FOUND}; }}; done",
        f"tar {flags} - {quoted}",
    ]
    return "; ".join(statements)