"""Offline contracts for strict SSH command construction."""

from __future__ import annotations

import shlex
import subprocess
from pathlib import Path
from unittest.mock import Mock

import pytest
import yaml

from obliteratus import __version__
from obliteratus.remote import RemoteConfig, RemoteRunner


def test_remote_commands_require_strict_host_key_verification(tmp_path):
    key = tmp_path / "key"
    known_hosts = tmp_path / "known_hosts"
    config = RemoteConfig(
        host="compute.example",
        user="runner",
        port=2222,
        ssh_key=str(key),
        known_hosts_file=str(known_hosts),
    )
    runner = RemoteRunner(config)
    for command in (runner._ssh_base_cmd(), runner._scp_base_cmd()):
        assert "StrictHostKeyChecking=yes" in command
        assert f"UserKnownHostsFile={known_hosts}" in command
        assert "StrictHostKeyChecking=no" not in command
        assert str(key) in command


def test_remote_config_accepts_versioned_known_hosts_setting():
    config = RemoteConfig.from_dict(
        {
            "host": "compute.example",
            "user": "runner",
            "known_hosts_file": "/secure/known_hosts",
            "unknown": "ignored",
        }
    )
    assert config.known_hosts_file == "/secure/known_hosts"
    assert config.ssh_target == "runner@compute.example"


def _completed(returncode=0, stdout="", stderr=""):
    return subprocess.CompletedProcess([], returncode, stdout=stdout, stderr=stderr)


def test_remote_config_cli_defaults_and_optional_command_flags_are_shell_safe():
    config = RemoteConfig.from_cli_args(
        "runner@compute.example",
        python="/opt/python builds/python",
        remote_dir="/srv/remote work",
        gpus="02, 0",
    )
    runner = RemoteRunner(config, on_log=lambda _message: None)
    command = runner.build_obliterate_command(
        "org/model; touch /tmp/injected",
        method="advanced",
        device="cuda",
        dtype="float16",
        quantization="4bit",
        n_directions=3,
        direction_method="svd",
        regularization=0.2,
        refinement_passes=2,
        large_model=True,
        verify_sample_size=7,
        min_layer_fraction=0.1,
        max_layer_fraction=0.9,
        harmless_pc_count=4,
        shield_concept_count=5,
        shield_ridge=0.05,
        shield_residualize=True,
        shield_layer_penalty=0.3,
        projection_target="attention",
        projection_row_fraction=0.25,
        refusal_max_tokens=256,
    )
    tokens = shlex.split(command)
    assert tokens[:4] == ["env", "CUDA_VISIBLE_DEVICES=2,0", "/opt/python builds/python", "-m"]
    assert tokens[4:7] == ["obliteratus", "obliterate", "org/model; touch /tmp/injected"]
    expected = {
        "--quantization": "4bit",
        "--n-directions": "3",
        "--direction-method": "svd",
        "--regularization": "0.2",
        "--refinement-passes": "2",
        "--verify-sample-size": "7",
        "--min-layer-fraction": "0.1",
        "--max-layer-fraction": "0.9",
        "--harmless-pc-count": "4",
        "--shield-concept-count": "5",
        "--shield-ridge": "0.05",
        "--shield-layer-penalty": "0.3",
        "--projection-target": "attention",
        "--projection-row-fraction": "0.25",
        "--refusal-max-tokens": "256",
    }
    for flag, value in expected.items():
        assert tokens[tokens.index(flag) + 1] == value
    assert "--large-model" in tokens
    assert "--shield-residualize" in tokens


def test_remote_config_rejects_invalid_install_timeout():
    with pytest.raises(ValueError, match="install timeout"):
        RemoteConfig(host="host", install_timeout=0)


def test_run_and_tourney_commands_quote_all_public_values():
    runner = RemoteRunner(
        RemoteConfig(host="host", user="runner", gpus="all"),
        on_log=lambda _message: None,
    )
    assert shlex.split(
        runner.build_run_command("/tmp/a config.yml", output_dir="/tmp/out dir", preset="x; echo bad")
    ) == [
        "python3", "-m", "obliteratus", "run", "/tmp/a config.yml",
        "--output-dir", "/tmp/out dir", "--preset", "x; echo bad",
    ]
    tokens = shlex.split(
        runner.build_tourney_command(
            "org/model", output_dir="/tmp/out dir", quantization="8bit",
            hub_org="org; bad", hub_repo="org/repo bad", methods=["basic", "advanced"],
            dataset="data; bad",
        )
    )
    assert tokens[tokens.index("--hub-org") + 1] == "org; bad"
    assert tokens[tokens.index("--hub-repo") + 1] == "org/repo bad"
    assert tokens[tokens.index("--dataset") + 1] == "data; bad"
    assert tokens[-3:] == ["--methods", "basic", "advanced"]
    assert shlex.split(runner.build_run_command("config.yml", preset="quick")) == [
        "python3", "-m", "obliteratus", "run", "config.yml", "--preset", "quick",
    ]


def test_run_ssh_non_stream_uses_argv_and_propagates_timeout(monkeypatch):
    observed = {}

    def fake_run(command, **kwargs):
        observed.update(command=command, kwargs=kwargs)
        return _completed(stdout="ok\n")

    monkeypatch.setattr(subprocess, "run", fake_run)
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
    result = runner.run_ssh("printf '%s' 'safe value'", timeout=9)
    assert result.stdout == "ok\n"
    assert observed["command"][-1] == "printf '%s' 'safe value'"
    assert observed["kwargs"] == {"capture_output": True, "text": True, "timeout": 9}


class _StreamProcess:
    def __init__(self, stdout):
        self.stdout = stdout
        self.returncode = 0
        self.killed = False
        self.wait_calls = 0

    def wait(self, timeout=None):
        self.wait_calls += 1
        return self.returncode

    def kill(self):
        self.killed = True


def test_run_ssh_streams_lines_and_returns_exit_status(monkeypatch):
    process = _StreamProcess(["first\n", "second\n"])
    monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
    assert runner.run_ssh("command", stream=True, timeout=4) == 0
    assert logs == ["first", "second"]
    assert process.wait_calls == 1


def test_run_ssh_timeout_kills_and_reaps_process(monkeypatch):
    process = _StreamProcess([])

    def wait(timeout=None):
        process.wait_calls += 1
        if process.wait_calls == 1:
            raise subprocess.TimeoutExpired("ssh", timeout)
        return 0

    process.wait = wait
    monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
    assert runner.run_ssh("command", stream=True, timeout=4) == 124
    assert process.killed is True
    assert process.wait_calls == 2
    assert any("timed out" in line for line in logs)


class _CancelledOutput:
    def __iter__(self):
        return self

    def __next__(self):
        raise KeyboardInterrupt


def test_run_ssh_cancellation_kills_reaps_and_reraises(monkeypatch):
    process = _StreamProcess(_CancelledOutput())
    monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
    with pytest.raises(KeyboardInterrupt):
        runner.run_ssh("command", stream=True)
    assert process.killed is True
    assert process.wait_calls == 1
    assert any("cancelled" in line for line in logs)


def test_run_ssh_rejects_malformed_process_without_stdout(monkeypatch):
    process = _StreamProcess(None)
    monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
    with pytest.raises(RuntimeError, match="stdout"):
        runner.run_ssh("command", stream=True)
    assert process.killed is True
    assert process.wait_calls == 1


def test_connection_and_gpu_probes_cover_success_and_malformed_responses():
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host", gpus="1"), on_log=logs.append)
    runner.run_ssh = Mock(return_value=_completed(stdout="ok\n"))
    assert runner.check_connection() is True
    runner.run_ssh.return_value = 0
    assert runner.check_connection() is False

    runner.run_ssh.return_value = _completed(stdout="0, A100, 80 GiB, 70 GiB\n1, A100, 80 GiB, 60 GiB\n")
    assert runner.check_gpu().splitlines() == [
        "0, A100, 80 GiB, 70 GiB",
        "1, A100, 80 GiB, 60 GiB",
    ]
    assert any("Selected GPUs: 1" in line for line in logs)
    runner.run_ssh.return_value = _completed(returncode=1, stderr="missing")
    assert runner.check_gpu() is None


def test_gpu_probe_reports_all_devices_when_unrestricted():
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host", gpus="all"), on_log=logs.append)
    runner.run_ssh = Mock(return_value=_completed(stdout="0, A100\n"))
    assert runner.check_gpu() == "0, A100"
    assert any("Using: all 1 GPUs" in line for line in logs)


def test_ensure_obliteratus_accepts_exact_version_or_installs_and_verifies():
    logs = []
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
    runner.run_ssh = Mock(return_value=_completed(stdout=f"{__version__}\n"))
    assert runner.ensure_obliteratus() is True
    assert runner.run_ssh.call_count == 1

    runner.run_ssh = Mock(side_effect=[_completed(stdout="0.0.1\n"), 0, _completed(stdout=f"{__version__}\n")])
    assert runner.ensure_obliteratus() is True
    install_command = runner.run_ssh.call_args_list[1].args[0]
    assert shlex.split(install_command)[-1] == "git+https://github.com/elder-plinius/OBLITERATUS.git"


def test_ensure_obliteratus_reports_install_and_verification_failures():
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
    runner.run_ssh = Mock(side_effect=[_completed(returncode=1), 1])
    assert runner.ensure_obliteratus() is False
    runner.run_ssh = Mock(side_effect=[_completed(stdout="old\n"), 0, _completed(stdout="still-old\n")])
    assert runner.ensure_obliteratus() is False


def test_result_sync_creates_local_directory_and_quotes_remote_path(monkeypatch, tmp_path):
    responses = iter([_completed(), _completed(returncode=1, stderr="denied")])
    observed = []

    def fake_run(command, **_kwargs):
        observed.append(command)
        return next(responses)

    monkeypatch.setattr(subprocess, "run", fake_run)
    runner = RemoteRunner(RemoteConfig(host="host", user="runner"), on_log=lambda _message: None)
    local = tmp_path / "local results"
    assert runner.sync_results_back("/tmp/remote results", str(local)) is True
    assert local.is_dir()
    assert observed[0][-2] == "runner@host:'/tmp/remote results/'"
    assert runner.sync_results_back("/tmp/remote results", str(local)) is False


def test_upload_config_returns_remote_path_or_raises(monkeypatch, tmp_path):
    config_path = tmp_path / "study config.yml"
    config_path.write_text("model: fixture\nremote:\n  host: compute.example\n")
    responses = iter([_completed(), _completed(returncode=1, stderr="denied")])
    observed = []
    uploaded_payloads = []

    def fake_run(command, **_kwargs):
        observed.append(command)
        uploaded_payloads.append(yaml.safe_load(Path(command[-2]).read_text()))
        return next(responses)

    monkeypatch.setattr(subprocess, "run", fake_run)
    runner = RemoteRunner(
        RemoteConfig(host="host", user="runner", remote_dir="/srv/remote work"),
        on_log=lambda _message: None,
    )
    runner.run_ssh = Mock()
    assert runner.upload_config(str(config_path)) == "/srv/remote work/config.yaml"
    assert observed[0][-1] == "runner@host:'/srv/remote work/config.yaml'"
    assert uploaded_payloads[0] == {"model": "fixture"}
    with pytest.raises(RuntimeError, match="denied"):
        runner.upload_config(str(config_path))


def test_upload_config_requires_yaml_mapping(monkeypatch, tmp_path):
    config_path = tmp_path / "study.yml"
    config_path.write_text("- not\n- a\n- mapping\n")
    runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
    runner.run_ssh = Mock()
    monkeypatch.setattr(subprocess, "run", Mock())
    with pytest.raises(ValueError, match="must contain a YAML mapping"):
        runner.upload_config(str(config_path))
    subprocess.run.assert_not_called()


def _prepared_runner(*, sync_results=True):
    runner = RemoteRunner(
        RemoteConfig(host="host", remote_dir="/srv/run", sync_results=sync_results),
        on_log=lambda _message: None,
    )
    runner.check_connection = Mock(return_value=True)
    runner.check_gpu = Mock(return_value="gpu")
    runner.ensure_obliteratus = Mock(return_value=True)
    return runner


def test_remote_obliterate_orchestration_success_failure_and_sync_paths():
    runner = _prepared_runner(sync_results=False)
    runner.run_ssh = Mock(side_effect=[_completed(), 0])
    assert runner.run_obliterate("org/model") == "/srv/run/output/org_model"

    runner = _prepared_runner()
    runner.run_ssh = Mock(side_effect=[_completed(), 9])
    assert runner.run_obliterate("org/model") is None

    runner = _prepared_runner()
    runner.run_ssh = Mock(side_effect=[_completed(), 0])
    runner.sync_results_back = Mock(return_value=True)
    assert runner.run_obliterate("org/model", local_output_dir="local") == "local"
    runner.run_ssh = Mock(side_effect=[_completed(), 0])
    runner.sync_results_back.return_value = False
    assert runner.run_obliterate("org/model") is None


def test_remote_obliterate_stops_at_connection_or_install_failure():
    runner = _prepared_runner()
    runner.check_connection.return_value = False
    assert runner.run_obliterate("model") is None
    runner.check_connection.return_value = True
    runner.ensure_obliteratus.return_value = False
    assert runner.run_obliterate("model") is None


def test_remote_config_orchestration_success_failure_and_sync_paths():
    runner = _prepared_runner(sync_results=False)
    runner.upload_config = Mock(return_value="/srv/run/config.yaml")
    runner.run_ssh = Mock(return_value=0)
    assert runner.run_config("local.yml", preset="quick") == "/srv/run/results"

    runner = _prepared_runner()
    runner.upload_config = Mock(return_value="/srv/run/config.yaml")
    runner.run_ssh = Mock(return_value=2)
    assert runner.run_config("local.yml") is None

    runner.run_ssh.return_value = 0
    runner.sync_results_back = Mock(return_value=True)
    assert runner.run_config("local.yml", local_output_dir="local") == "local"
    runner.sync_results_back.return_value = False
    assert runner.run_config("local.yml") is None


def test_remote_config_stops_at_connection_or_install_failure():
    runner = _prepared_runner()
    runner.check_connection.return_value = False
    assert runner.run_config("local.yml") is None
    runner.check_connection.return_value = True
    runner.ensure_obliteratus.return_value = False
    assert runner.run_config("local.yml") is None


def test_remote_tourney_orchestration_success_failure_and_sync_paths():
    runner = _prepared_runner(sync_results=False)
    runner.run_ssh = Mock(return_value=0)
    assert runner.run_tourney("org/model") == "/srv/run/tourney/org_model"

    runner = _prepared_runner()
    runner.run_ssh = Mock(return_value=3)
    assert runner.run_tourney("org/model") is None

    runner.run_ssh.return_value = 0
    runner.sync_results_back = Mock(return_value=True)
    assert runner.run_tourney("org/model", local_output_dir="local") == "local"
    runner.sync_results_back.return_value = False
    assert runner.run_tourney("org/model") is None


def test_remote_tourney_stops_at_connection_or_install_failure():
    runner = _prepared_runner()
    runner.check_connection.return_value = False
    assert runner.run_tourney("model") is None
    runner.check_connection.return_value = True
    runner.ensure_obliteratus.return_value = False
    assert runner.run_tourney("model") is None