"""tias-remote CLI のテスト（SPEC_03）．ssh はすべてモックし，実サーバに接続しない.

- 接続できないサーバがあっても他のサーバで続行する（全台不可ならエラー）
- `TIAS_REMOTE_JUMP` が空文字なら踏み台を使わない
- 手元のホスト名がサーバ名と一致したら止まる（`TIAS_REMOTE_ALLOW_SELF=1` で無効化）
- pull は回収ホワイトリストに一致するファイルだけを持ち帰り，回収しなかったファイルの名前を
  出さない．受け取った tar に不正なメンバーがあれば何も展開しない
"""

from __future__ import annotations

import io
import logging
import tarfile
import threading
from collections.abc import Callable
from pathlib import Path

import pytest

import tias_judge_model.cli.remote as cli
from tias_judge_model.remote.commands import (
    EXIT_NOT_FOUND,
    PROBE_APPS_MARKER,
    PROBE_APPS_RC_PREFIX,
    PROBE_END_MARKER,
    PROBE_GPU_RC_PREFIX,
    PROBE_GPUS_MARKER,
    PROBE_OWNERS_MARKER,
    PROBE_SAMPLES_MARKER,
    PROBE_SESSIONS_MARKER,
)
from tias_judge_model.remote.gpu import GpuState

NEWTON_ADDR = "192.168.23.14"
ARCHIMEDES_ADDR = "192.168.23.15"
RUN_ID = "20260927-120000_newton_abc1234"

CONFIG_YAML = f"""\
jump_host: labpc
connect_timeout_sec: 15
gpu_free:
  max_memory_used_mib: 1024
  max_utilization_percent: 10
  observe_sec: 10
  sample_interval_ms: 500
servers:
  - name: Newton
    address: {NEWTON_ADDR}
    user: hasegawa
    priority: 1
    remote_dir: ~/work/TIASJudgeModel
    data_dir: ~/data/TIASJudgeModel
    uv: ~/.local/bin/uv
    data_ready: true
  - name: Archimedes
    address: {ARCHIMEDES_ADDR}
    user: hasegawa
    priority: 2
    remote_dir: ~/work/TIASJudgeModel
    data_dir: ~/data/TIASJudgeModel
    uv: ~/.local/bin/uv
    data_ready: true
"""


def _probe_stdout(gpus: list[tuple[int, int, int]], apps: list[str] | None = None) -> bytes:
    """probe の出力（gpus は (index, 使用中メモリ, 使用率) のリスト．観測サンプルも同じ値）."""
    lines = [PROBE_SESSIONS_MARKER, PROBE_GPUS_MARKER]
    lines += [f"{i}, GPU-{i}, NVIDIA RTX A6000, {mem}, 49140, {util}" for i, mem, util in gpus]
    lines.append(f"{PROBE_GPU_RC_PREFIX}0")
    lines.append(PROBE_SAMPLES_MARKER)
    lines += [f"{i}, {util}, {mem}" for i, mem, util in gpus for _ in range(3)]
    lines.append(PROBE_APPS_MARKER)
    lines += apps or []
    lines += [f"{PROBE_APPS_RC_PREFIX}0", PROBE_OWNERS_MARKER, PROBE_END_MARKER]
    return ("\n".join(lines) + "\n").encode()


class FakeSsh:
    """`_run_ssh` の代わり．接続先アドレスごとに応答を返し，呼び出しを記録する."""

    def __init__(self, responses: dict[str, Callable[[list[str]], cli._SshResult]]) -> None:
        self.responses = responses
        self.calls: list[list[str]] = []
        self._lock = threading.Lock()

    def __call__(
        self, argv: list[str], *, input_bytes: bytes | None = None, timeout: float | None = None
    ) -> cli._SshResult:
        with self._lock:
            self.calls.append(list(argv))
        for address, respond in self.responses.items():
            if f"hasegawa@{address}" in argv:
                return respond(argv)
        raise AssertionError("想定外の接続先")


def _ok(stdout: bytes) -> Callable[[list[str]], cli._SshResult]:
    return lambda argv: cli._SshResult(0, stdout, b"")


def _fail(rc: int = 255, stderr: bytes = b"ssh: connect to host: Connection timed out"):
    return lambda argv: cli._SshResult(rc, b"", stderr)


@pytest.fixture
def project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    """手元のプロジェクトルートを模したディレクトリ（ssh は一切起動させない）."""
    root = tmp_path / "proj"
    (root / "src" / "tias_judge_model").mkdir(parents=True)
    (root / "src" / "tias_judge_model" / "__init__.py").write_bytes(b"")
    (root / "pyproject.toml").write_bytes(b"[project]\nname = 'x'\n")
    (root / "configs").mkdir()
    (root / "configs" / "remote.yaml").write_text(CONFIG_YAML, encoding="utf-8")
    monkeypatch.chdir(root)
    for var in ("TIAS_REMOTE_JUMP", "TIAS_REMOTE_ALLOW_SELF"):
        monkeypatch.delenv(var, raising=False)
    monkeypatch.setenv("TIAS_REMOTE_SSH", "ssh")
    monkeypatch.setattr(cli.socket, "gethostname", lambda: "my-laptop")

    def _no_process(*args: object, **kwargs: object) -> None:
        raise AssertionError("テストが実プロセス（ssh）を起動しようとした")

    monkeypatch.setattr(cli, "_popen", _no_process)
    monkeypatch.setattr(cli.subprocess, "Popen", _no_process)
    return root


def _spy_selection(monkeypatch: pytest.MonkeyPatch) -> list:
    picked: list = []
    original = cli.select_gpu

    def spy(*args, **kwargs):
        result = original(*args, **kwargs)
        picked.append(result)
        return result

    monkeypatch.setattr(cli, "select_gpu", spy)
    return picked


# ---------------------------------------------------------------------------
# 接続できないサーバ
# ---------------------------------------------------------------------------


class TestUnreachableServers:
    def test_gpus_continues_with_other_server(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        picked = _spy_selection(monkeypatch)
        assert cli.main(["gpus"]) == 0
        assert picked[0] is not None
        assert (picked[0].server.name, picked[0].gpu_index) == ("Archimedes", 0)
        assert len(fake.calls) == 2

    def test_run_selects_gpu_on_reachable_server(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        picked = _spy_selection(monkeypatch)
        rc = cli.main(["run", "--name", "exp1", "--dry-run", "--", "python", "-V"])
        assert rc == 0
        assert picked[0].server.name == "Archimedes"
        assert picked[0].state is GpuState.UNUSED

    def test_unparsable_probe_is_treated_as_unreachable(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh(
            {NEWTON_ADDR: _ok(b"garbage\n"), ARCHIMEDES_ADDR: _ok(_probe_stdout([(1, 5, 0)]))}
        )
        monkeypatch.setattr(cli, "_run_ssh", fake)
        picked = _spy_selection(monkeypatch)
        assert cli.main(["gpus"]) == 0
        assert (picked[0].server.name, picked[0].gpu_index) == ("Archimedes", 1)

    @pytest.mark.parametrize("subcommand", [["gpus"], ["status"], ["push"]])
    def test_all_unreachable_is_error(
        self, project: Path, monkeypatch: pytest.MonkeyPatch, subcommand: list[str]
    ) -> None:
        fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _fail()})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(subcommand) != 0

    def test_run_all_unreachable_is_error(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _fail()})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["run", "--name", "exp1", "--dry-run", "--", "python"]) != 0


class TestRunSelection:
    def test_unused_on_lower_priority_server_is_chosen(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        newton = _probe_stdout([(0, 300, 0)], apps=["GPU-0, 4321, 300"])
        fake = FakeSsh({NEWTON_ADDR: _ok(newton), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        picked = _spy_selection(monkeypatch)
        assert cli.main(["run", "--name", "exp1", "--dry-run", "--", "python"]) == 0
        assert picked[0].server.name == "Archimedes"

    def test_no_free_gpu_does_not_start_job(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        busy = _probe_stdout([(0, 20000, 95)])
        fake = FakeSsh({NEWTON_ADDR: _ok(busy), ARCHIMEDES_ADDR: _ok(busy)})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["run", "--name", "exp1", "--", "python"]) != 0
        # probe の 2 回だけで，起動のための ssh は呼ばない
        assert len(fake.calls) == 2

    def test_gpu_without_server_is_rejected_before_ssh(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh({})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["run", "--name", "exp1", "--gpu", "0", "--", "python"]) != 0
        assert fake.calls == []

    def test_explicit_busy_gpu_is_error_without_force(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        busy = _probe_stdout([(0, 20000, 95)])
        fake = FakeSsh({NEWTON_ADDR: _ok(busy), ARCHIMEDES_ADDR: _ok(busy)})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        args = ["run", "--name", "e", "--server", "Newton", "--gpu", "0", "--dry-run", "--", "x"]
        assert cli.main(args) != 0
        assert cli.main([*args[:7], "--force", *args[7:]]) == 0


# ---------------------------------------------------------------------------
# 環境変数
# ---------------------------------------------------------------------------


class TestEnvironment:
    def _run_gpus(self, monkeypatch: pytest.MonkeyPatch) -> FakeSsh:
        fake = FakeSsh(
            {
                NEWTON_ADDR: _ok(_probe_stdout([(0, 5, 0)])),
                ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)])),
            }
        )
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["gpus"]) == 0
        return fake

    def test_config_jump_host_is_used(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        fake = self._run_gpus(monkeypatch)
        for argv in fake.calls:
            assert argv[argv.index("-J") + 1] == "labpc"
            assert "BatchMode=yes" in argv

    def test_empty_jump_env_means_direct(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        monkeypatch.setenv("TIAS_REMOTE_JUMP", "")
        fake = self._run_gpus(monkeypatch)
        assert fake.calls
        for argv in fake.calls:
            assert "-J" not in argv

    def test_self_server_is_detected(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(cli.socket, "gethostname", lambda: "Newton")
        fake = FakeSsh({})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["gpus"]) != 0
        assert fake.calls == []

    def test_allow_self_disables_detection(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        monkeypatch.setattr(cli.socket, "gethostname", lambda: "Newton")
        monkeypatch.setenv("TIAS_REMOTE_ALLOW_SELF", "1")
        fake = self._run_gpus(monkeypatch)
        assert len(fake.calls) == 2


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

LISTING_NAMES = [
    "config.yaml",
    "history.csv",
    "metrics.json",
    "confusion_matrix.csv",
    "log.txt",
    "figures/confusion_matrix.png",
    "split.json",
    "predictions.csv",
    "P0001_yamada.png",
    "manifest.csv",
    "checkpoints/best.pt",
]
ALLOWED = LISTING_NAMES[:6]
SECRET_NAMES = ["split.json", "predictions.csv", "P0001_yamada", "manifest.csv"]


def _listing() -> bytes:
    return b"".join(f"f {RUN_ID}/{name}\0".encode() for name in LISTING_NAMES)


def _tar_of(names: list[str]) -> bytes:
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w", format=tarfile.PAX_FORMAT) as tar:
        for name in names:
            data = f"content of {name}".encode()
            info = tarfile.TarInfo(name)
            info.size = len(data)
            tar.addfile(info, io.BytesIO(data))
    return buffer.getvalue()


class FakeReceive:
    """`_receive_tar` の代わり．要求されたパスの tar（と任意の追加メンバー）を返す."""

    def __init__(self, extra: list[str] | None = None) -> None:
        self.extra = extra or []
        self.requested: list[str] = []

    def __call__(self, argv: list[str], payload: bytes, destination) -> cli._SshResult:
        self.requested = [p.decode() for p in payload.split(b"\0") if p]
        destination.write(_tar_of(self.requested + self.extra))
        return cli._SshResult(0, b"", b"")


def _setup_pull(monkeypatch: pytest.MonkeyPatch, receive: FakeReceive) -> FakeSsh:
    fake = FakeSsh(
        {
            NEWTON_ADDR: _ok(_listing()),
            ARCHIMEDES_ADDR: _fail(rc=EXIT_NOT_FOUND, stderr=b""),
        }
    )
    monkeypatch.setattr(cli, "_run_ssh", fake)
    monkeypatch.setattr(cli, "_receive_tar", receive)
    return fake


def _local_files(root: Path) -> list[str]:
    run = root / "runs" / RUN_ID
    if not run.exists():
        return []
    return sorted(p.relative_to(run).as_posix() for p in run.rglob("*") if p.is_file())


class TestPull:
    def test_pulls_only_whitelisted_files(
        self,
        project: Path,
        monkeypatch: pytest.MonkeyPatch,
        capsys: pytest.CaptureFixture[str],
        caplog: pytest.LogCaptureFixture,
    ) -> None:
        caplog.set_level(logging.DEBUG)
        receive = FakeReceive()
        _setup_pull(monkeypatch, receive)
        assert cli.main(["pull", RUN_ID]) == 0
        assert sorted(receive.requested) == sorted(f"{RUN_ID}/{n}" for n in ALLOWED)
        assert _local_files(project) == sorted(ALLOWED)
        captured = capsys.readouterr()
        output = captured.out + captured.err + caplog.text
        for secret in [*SECRET_NAMES, "best.pt"]:
            assert secret not in output
        # 回収しなかった件数（split・predictions・画像・manifest・checkpoint の 5 件）は示す
        assert "5 件" in output

    def test_with_checkpoints(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        receive = FakeReceive()
        _setup_pull(monkeypatch, receive)
        assert cli.main(["pull", RUN_ID, "--with-checkpoints"]) == 0
        assert _local_files(project) == sorted([*ALLOWED, "checkpoints/best.pt"])

    @pytest.mark.parametrize(
        "extra",
        [
            f"{RUN_ID}/split.json",
            f"{RUN_ID}/P0001_yamada.png",
            f"{RUN_ID}/../evil.txt",
            "../evil.txt",
            "other_run/config.yaml",
        ],
    )
    def test_malicious_archive_extracts_nothing(
        self,
        project: Path,
        monkeypatch: pytest.MonkeyPatch,
        capsys: pytest.CaptureFixture[str],
        caplog: pytest.LogCaptureFixture,
        extra: str,
    ) -> None:
        caplog.set_level(logging.DEBUG)
        _setup_pull(monkeypatch, FakeReceive(extra=[extra]))
        assert cli.main(["pull", RUN_ID]) != 0
        assert _local_files(project) == []
        assert not (project / "runs" / "evil.txt").exists()
        assert not (project / "evil.txt").exists()
        captured = capsys.readouterr()
        output = captured.out + captured.err + caplog.text
        assert "P0001_yamada" not in output
        assert "split.json" not in output

    @pytest.mark.parametrize("run_id", ["../etc", "a/b", "a\\b", "..", "x..y"])
    def test_invalid_run_id_is_rejected_before_ssh(
        self, project: Path, monkeypatch: pytest.MonkeyPatch, run_id: str
    ) -> None:
        fake = FakeSsh({})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["pull", run_id]) != 0
        assert fake.calls == []

    def test_run_on_multiple_servers_is_error(
        self, project: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        fake = FakeSsh({NEWTON_ADDR: _ok(_listing()), ARCHIMEDES_ADDR: _ok(_listing())})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        monkeypatch.setattr(cli, "_receive_tar", FakeReceive())
        assert cli.main(["pull", RUN_ID]) != 0
        assert cli.main(["pull", RUN_ID, "--server", "Newton"]) == 0

    def test_run_not_found_is_error(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        missing = _fail(rc=EXIT_NOT_FOUND, stderr=b"")
        fake = FakeSsh({NEWTON_ADDR: missing, ARCHIMEDES_ADDR: missing})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["pull", RUN_ID]) != 0


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


class TestPush:
    def test_dry_run_lists_code_but_not_excluded(
        self, project: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
    ) -> None:
        for rel in ["data/synthetic/manifest.csv", "runs/x/metrics.json", ".claude/a.md"]:
            path = project / rel
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(b"x")
        fake = FakeSsh({})
        monkeypatch.setattr(cli, "_run_ssh", fake)
        assert cli.main(["push", "--dry-run"]) == 0
        out = capsys.readouterr().out
        assert "pyproject.toml" in out
        assert "configs/remote.yaml" in out
        assert "data/synthetic" not in out
        assert "runs/x" not in out
        assert ".claude" not in out
        assert fake.calls == []