"""ssh 引数とサーバ上で実行するコマンドの組み立てのテスト（SPEC_03「接続」「取得」「run」）."""

from __future__ import annotations

import pytest

from tias_judge_model.remote.commands import (
    build_job_env,
    build_job_script,
    build_probe_command,
    build_ssh_argv,
    format_git_commit,
    validate_job_name,
)


def _argv(jump: str | None, timeout: int = 15) -> list[str]:
    return build_ssh_argv(
        "ssh",
        user="hasegawa",
        address="192.168.23.14",
        jump_host=jump,
        connect_timeout_sec=timeout,
        remote_command="echo hi",
    )


def _has_option(argv: list[str], value: str) -> bool:
    return any(argv[i] == "-o" and argv[i + 1] == value for i in range(len(argv) - 1))


class TestBuildSshArgv:
    def test_with_jump_host(self) -> None:
        argv = _argv("labpc")
        assert argv[0] == "ssh"
        j = argv.index("-J")
        assert argv[j + 1] == "labpc"
        assert "hasegawa@192.168.23.14" in argv
        assert argv[-1] == "echo hi"
        # 接続先より前にオプションがある
        assert j < argv.index("hasegawa@192.168.23.14")

    def test_without_jump_host(self) -> None:
        argv = _argv(None)
        assert "-J" not in argv
        assert "hasegawa@192.168.23.14" in argv

    @pytest.mark.parametrize("jump", ["labpc", None])
    def test_batch_mode_and_timeout_always(self, jump: str | None) -> None:
        argv = _argv(jump, timeout=7)
        assert _has_option(argv, "BatchMode=yes")
        assert _has_option(argv, "ConnectTimeout=7")

    def test_remote_command_is_single_argument(self) -> None:
        argv = build_ssh_argv(
            "ssh",
            user="u",
            address="h",
            jump_host=None,
            connect_timeout_sec=5,
            remote_command="cd ~/x; echo a b",
        )
        assert argv[-1] == "cd ~/x; echo a b"

    @pytest.mark.parametrize("bad", ["-oProxyCommand=evil", "a b", ""])
    def test_rejects_option_like_destination(self, bad: str) -> None:
        with pytest.raises(ValueError):
            build_ssh_argv(
                "ssh",
                user="u",
                address=bad,
                jump_host=None,
                connect_timeout_sec=5,
                remote_command="true",
            )


class TestBuildProbeCommand:
    def test_queries_listed_in_spec(self) -> None:
        command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
        assert "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu" in command
        assert "--query-compute-apps=gpu_uuid,pid,used_memory" in command
        assert "--format=csv,noheader,nounits" in command

    def test_observation_uses_interval_and_duration(self) -> None:
        command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
        assert "-lms 500" in command
        assert "timeout 10" in command
        assert "pmon -s u" in command

    def test_observe_zero_does_not_observe(self) -> None:
        command = build_probe_command("~/work/x", observe_sec=0, sample_interval_ms=500)
        assert "-lms" not in command
        assert "pmon" not in command
        assert "--query-gpu=index,uuid,name" in command

    def test_is_single_line(self) -> None:
        command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
        assert "\n" not in command


class TestJobSettings:
    @pytest.mark.parametrize("name", ["a", "exp1", "Exp_1-b", "x" * 64])
    def test_valid_job_names(self, name: str) -> None:
        assert validate_job_name(name) == name

    @pytest.mark.parametrize("name", ["", "_a", "-a", "a b", "a.b", "a:b", "x" * 65, "a/b"])
    def test_invalid_job_names(self, name: str) -> None:
        with pytest.raises(ValueError):
            validate_job_name(name)

    def test_job_env_contains_selected_gpu_and_extra(self) -> None:
        env = build_job_env(2, ["FOO=bar", "EMPTY="])
        assert env["CUDA_VISIBLE_DEVICES"] == "2"
        assert env["FOO"] == "bar"
        assert env["EMPTY"] == ""


class TestGitCommit:
    """`TIAS_GIT_COMMIT`（SPEC_03「run」・SPEC_01「実験の出力」）."""

    def test_clean_tree_passes_short_hash(self) -> None:
        # git rev-parse --short HEAD の出力（末尾改行付き）
        assert format_git_commit("abc1234\n", dirty=False) == "abc1234"

    def test_dirty_tree_appends_suffix(self) -> None:
        assert format_git_commit("abc1234\n", dirty=True) == "abc1234-dirty"

    def test_job_env_contains_git_commit(self) -> None:
        env = build_job_env(0, [], git_commit="abc1234-dirty")
        assert env["TIAS_GIT_COMMIT"] == "abc1234-dirty"

    def test_job_env_without_git_commit_does_not_set_it(self) -> None:
        env = build_job_env(0, [])
        assert "TIAS_GIT_COMMIT" not in env

    @pytest.mark.parametrize("git_commit", ["abc1234", None])
    def test_env_option_cannot_override_git_commit(self, git_commit: str | None) -> None:
        # git が使えず渡さないときも --env で指定させない（実験の記録を迂回させないため）
        with pytest.raises(ValueError):
            build_job_env(0, ["TIAS_GIT_COMMIT=fake"], git_commit=git_commit)


class TestReservedEnv:
    @pytest.mark.parametrize(
        "key",
        [
            "CUDA_VISIBLE_DEVICES",
            "TIAS_DATA_ROOT",
            "CUDA_DEVICE_ORDER",
            "PYTHONUNBUFFERED",
            "TIAS_GIT_COMMIT",
        ],
    )
    @pytest.mark.parametrize("gpu_index", [0, None])
    def test_reserved_keys_cannot_be_overridden(self, key: str, gpu_index: int | None) -> None:
        with pytest.raises(ValueError):
            build_job_env(gpu_index, [f"{key}=x"], git_commit="abc1234")

    def test_fixed_values(self) -> None:
        env = build_job_env(3, [], git_commit="abc1234")
        assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
        assert env["PYTHONUNBUFFERED"] == "1"
        assert env["CUDA_VISIBLE_DEVICES"] == "3"


class TestCpuJob:
    """`run --cpu`（SPEC_03「run」）.

    CUDA_VISIBLE_DEVICES を空文字で渡し，ログの開始行の GPU は「なし（CPU）」とする．
    """

    def test_cpu_job_env_hides_gpus(self) -> None:
        env = build_job_env(None, ["FOO=bar"], git_commit="abc1234")
        assert "CUDA_VISIBLE_DEVICES" in env
        assert env["CUDA_VISIBLE_DEVICES"] == ""
        assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
        assert env["TIAS_GIT_COMMIT"] == "abc1234"
        assert env["FOO"] == "bar"

    def _script(self, gpu_index: int | None) -> str:
        return build_job_script(
            "~/work/TIASJudgeModel",
            "survey",
            ["tias-dataset", "survey", "--config", "configs/newton.yaml"],
            uv_path="~/.local/bin/uv",
            server_name="Newton",
            gpu_index=gpu_index,
            env=build_job_env(gpu_index, [], git_commit="abc1234"),
            path_env={"TIAS_DATA_ROOT": "~/data/TIASJudgeModel"},
        )

    def test_cpu_job_log_header_says_no_gpu(self) -> None:
        script = self._script(None)
        assert "gpu=なし（CPU）" in script
        assert "gpu=None" not in script
        # 空文字として export する（未設定ではなく空にして GPU を見えなくする）
        assert "export CUDA_VISIBLE_DEVICES=''" in script

    def test_gpu_job_log_header_has_index(self) -> None:
        script = self._script(1)
        assert "gpu=1" in script
        assert "なし（CPU）" not in script
        assert "export CUDA_VISIBLE_DEVICES=1" in script
        assert "export TIAS_GIT_COMMIT=abc1234" in script