"""Tests for client setup configuration"""

import shlex

import pytest

testinfra_hosts = [
    "borg-client",
    "borg-client-2",
    "borg-client-multi",
    "borg-client-nonroot",
    "borg-client-multi-keys",
    "borg-client-transition",
]


CLIENT_USER_MAP = {
    "borg-client": "root",
    "borg-client-2": "root",
    "borg-client-multi": "root",
    "borg-client-nonroot": "backupuser",
    "borg-client-multi-keys": "root",
    "borg-client-transition": "root",
}

CLIENT_SSH_KEY_TYPE_MAP = {
    "borg-client": "rsa",
    "borg-client-2": "ed25519",
    "borg-client-multi": "rsa",
    "borg-client-nonroot": "rsa",
    "borg-client-multi-keys": "ed25519",
    "borg-client-transition": "rsa",
}

CLIENT_SSH_KEY_PER_REPO_MAP = {
    "borg-client": False,
    "borg-client-2": False,
    "borg-client-multi": False,
    "borg-client-nonroot": False,
    "borg-client-multi-keys": True,
    "borg-client-transition": False,
}


def get_client_user(host):
    hostname = host.backend.get_hostname()
    return CLIENT_USER_MAP.get(hostname, "root")


def get_client_ssh_key_type(host):
    hostname = host.backend.get_hostname()
    return CLIENT_SSH_KEY_TYPE_MAP.get(hostname, "rsa")


def get_client_ssh_key_per_repo(host):
    hostname = host.backend.get_hostname()
    return CLIENT_SSH_KEY_PER_REPO_MAP.get(hostname, False)


def get_client_home(host):
    user = get_client_user(host)
    return f"/home/{user}" if user != "root" else "/root"


def get_backup_script_paths(hostname):
    if hostname == "borg-client-multi":
        return [
            "/usr/local/bin/run_borg_backup@configs",
            "/usr/local/bin/run_borg_backup@home-data",
        ]
    if hostname == "borg-client-multi-keys":
        return [
            "/usr/local/bin/run_borg_backup@configs-keys",
            "/usr/local/bin/run_borg_backup@home-data-keys",
        ]
    if hostname == "borg-client-nonroot":
        return ["/usr/local/bin/run_borg_backup@borg-server"]
    if hostname == "borg-client":
        return ["/usr/local/bin/run_borg_backup@borg-server"]
    if hostname == "borg-client-2":
        return ["/usr/local/bin/run_borg_backup@borg-server-2"]
    if hostname == "borg-client-transition":
        return ["/usr/local/bin/run_borg_backup@transition-repo"]
    return []


def get_prune_script_paths(hostname):
    if hostname == "borg-client":
        return ["/usr/local/bin/run_borg_prune@borg-server"]
    if hostname == "borg-client-2":
        return ["/usr/local/bin/run_borg_prune@borg-server-2"]
    if hostname == "borg-client-transition":
        return ["/usr/local/bin/run_borg_prune@transition-repo"]
    return []


def run_script_with_fake_borg(host, script_path, failed_command, failed_status):
    """Run a generated script with a fake borg and return its result and calls."""
    temp_dir_result = host.run("mktemp -d /tmp/borg-exit-test.XXXXXX")
    assert temp_dir_result.rc == 0
    temp_dir = temp_dir_result.stdout.strip()
    fake_borg_path = f"{temp_dir}/borg"
    log_path = f"{temp_dir}/calls"

    fake_borg = """#!/bin/bash
printf '%s\\n' "$1" >> "${BORG_TEST_LOG}"
if [ "$1" = "${BORG_FAIL_COMMAND}" ]; then
  exit "${BORG_FAIL_STATUS}"
fi
exit 0
"""
    setup = host.run(
        f"cat > {shlex.quote(fake_borg_path)} <<'EOF'\n"
        f"{fake_borg}"
        "EOF\n"
        f"chmod 0755 {shlex.quote(fake_borg_path)}"
    )
    assert setup.rc == 0

    try:
        result = host.run(
            "env "
            f"BORG_FAIL_COMMAND={shlex.quote(failed_command)} "
            f"BORG_FAIL_STATUS={failed_status} "
            f"BORG_TEST_LOG={shlex.quote(log_path)} "
            f"PATH={shlex.quote(temp_dir)}:$PATH "
            f"{shlex.quote(script_path)}"
        )
        calls = host.file(log_path).content_string.splitlines()
        return result, calls
    finally:
        host.run(f"rm -rf {shlex.quote(temp_dir)}")


class TestSSHSetup:
    def test_ssh_directory_exists(self, host):
        client_home = get_client_home(host)
        client_user = get_client_user(host)
        ssh_dir = host.file(f"{client_home}/.ssh")
        assert ssh_dir.exists
        assert ssh_dir.is_directory
        assert ssh_dir.user == client_user
        assert ssh_dir.group == client_user
        assert ssh_dir.mode == 0o700

    def test_ssh_private_key_exists(self, host):
        client_home = get_client_home(host)
        client_user = get_client_user(host)
        key_type = get_client_ssh_key_type(host)
        per_repo = get_client_ssh_key_per_repo(host)
        hostname = host.backend.get_hostname()

        if per_repo and hostname == "borg-client-multi-keys":
            key1 = host.file(
                f"{client_home}/.ssh/id_{key_type}_borgbackup_borg_server_configs_keys"
            )
            key2 = host.file(
                f"{client_home}/.ssh/id_{key_type}_borgbackup_borg_server_home_data_keys"
            )
            assert key1.exists
            assert key1.user == client_user
            assert key1.group == client_user
            assert key1.mode == 0o600
            assert key2.exists
            assert key2.user == client_user
            assert key2.group == client_user
            assert key2.mode == 0o600
        else:
            key = host.file(f"{client_home}/.ssh/id_{key_type}")
            assert key.exists
            assert key.user == client_user
            assert key.group == client_user
            assert key.mode == 0o600

    def test_ssh_public_key_exists(self, host):
        client_home = get_client_home(host)
        client_user = get_client_user(host)
        key_type = get_client_ssh_key_type(host)
        per_repo = get_client_ssh_key_per_repo(host)
        hostname = host.backend.get_hostname()

        if per_repo and hostname == "borg-client-multi-keys":
            key1 = host.file(
                f"{client_home}/.ssh/id_{key_type}_borgbackup_borg_server_configs_keys.pub"
            )
            key2 = host.file(
                f"{client_home}/.ssh/id_{key_type}_borgbackup_borg_server_home_data_keys.pub"
            )
            assert key1.exists
            assert key1.user == client_user
            assert key1.group == client_user
            assert key2.exists
            assert key2.user == client_user
            assert key2.group == client_user
        else:
            key = host.file(f"{client_home}/.ssh/id_{key_type}.pub")
            assert key.exists
            assert key.user == client_user
            assert key.group == client_user

    def test_known_hosts_contains_borg_server(self, host):
        client_home = get_client_home(host)
        known_hosts = host.file(f"{client_home}/.ssh/known_hosts")
        assert known_hosts.exists
        assert known_hosts.contains("borg-server")


class TestBackupScript:
    def test_backup_script_exists(self, host):
        hostname = host.backend.get_hostname()
        client_user = get_client_user(host)

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert script1.exists
            assert script1.user == client_user
            assert script1.group == client_user
            assert script1.mode == 0o711
            assert script2.exists
            assert script2.user == client_user
            assert script2.group == client_user
            assert script2.mode == 0o711
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert script1.exists
            assert script1.user == client_user
            assert script1.group == client_user
            assert script1.mode == 0o711
            assert script2.exists
            assert script2.user == client_user
            assert script2.group == client_user
            assert script2.mode == 0o711
        elif hostname == "borg-client-nonroot":
            script = host.file("/usr/local/bin/run_borg_backup@borg-server")
            assert script.exists
            assert script.user == client_user
            assert script.group == client_user
            assert script.mode == 0o711
        elif hostname in ("borg-client", "borg-client-2", "borg-client-transition"):
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.exists
            assert script.user == client_user
            assert script.group == client_user
            assert script.mode == 0o711
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")

    def test_backup_script_contains_borg_command(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert script1.contains("borg create")
            assert script2.contains("borg create")
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert script1.contains("borg create")
            assert script2.contains("borg create")
        elif hostname == "borg-client-nonroot":
            script = host.file("/usr/local/bin/run_borg_backup@borg-server")
            assert script.contains("borg create")
        elif hostname in ("borg-client", "borg-client-2", "borg-client-transition"):
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.contains("borg create")
            assert not script.contains("borg prune")
            assert not script.contains("borg compact")
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")


class TestPruneScript:
    def test_prune_script_created_when_pruning_enabled(self, host):
        hostname = host.backend.get_hostname()
        if hostname not in ("borg-client", "borg-client-2", "borg-client-transition"):
            return

        for script_path in get_prune_script_paths(hostname):
            script = host.file(script_path)
            assert script.exists
            assert script.mode == 0o711
            assert script.contains("borg prune")
            assert script.contains("borg compact")

    def test_scripts_have_valid_shell_syntax(self, host):
        hostname = host.backend.get_hostname()

        for script_path in get_backup_script_paths(hostname) + get_prune_script_paths(
            hostname
        ):
            result = host.run(f"bash -n {script_path}")
            assert result.rc == 0
            assert "ANSIBLE MANAGED BLOCK" not in host.file(script_path).content_string

    def test_after_backup_prune_script_flow_and_options(self, host):
        hostname = host.backend.get_hostname()
        if hostname != "borg-client":
            return

        backup = host.file("/usr/local/bin/run_borg_backup@borg-server")
        prune = host.file("/usr/local/bin/run_borg_prune@borg-server")
        content = prune.content_string

        assert "borg create" in backup.content_string
        assert "borg prune" not in backup.content_string
        assert "borg compact" not in backup.content_string

        prune_index = content.index("borg prune")
        prune_status_index = content.index("borg_prune_exit=$?")
        prune_guard_index = content.index('if [ "${borg_prune_exit}" -ne 0 ]; then')
        compact_index = content.index("borg compact")

        assert prune_index < prune_status_index < prune_guard_index < compact_index
        assert "--checkpoint-interval 1800" in content
        assert "--glob-archives '{hostname}-*'" in content
        assert "--keep-last 2" in content
        assert "--stats" in content
        assert "--threshold 1" in content

    def test_timer_prune_script_flow_and_options(self, host):
        hostname = host.backend.get_hostname()
        if hostname != "borg-client-2":
            return

        backup_script = host.file("/usr/local/bin/run_borg_backup@borg-server-2")
        prune_script = host.file("/usr/local/bin/run_borg_prune@borg-server-2")
        prune_content = prune_script.content_string

        assert "borg prune" not in backup_script.content_string
        assert "borg compact" not in backup_script.content_string

        prune_index = prune_content.index("borg prune")
        prune_status_index = prune_content.index("borg_prune_exit=$?")
        prune_guard_index = prune_content.index(
            'if [ "${borg_prune_exit}" -ne 0 ]; then'
        )
        compact_index = prune_content.index("borg compact")

        assert prune_index < prune_status_index < prune_guard_index < compact_index
        assert "--checkpoint-interval 1800" in prune_content
        assert "--glob-archives '{hostname}-*'" in prune_content
        assert "--keep-daily 7" in prune_content
        assert "--keep-weekly 4" in prune_content
        assert "--stats" in prune_content
        assert "--threshold 10" in prune_content

    def test_transition_host_uses_separate_prune_script(self, host):
        hostname = host.backend.get_hostname()
        if hostname != "borg-client-transition":
            return

        backup = host.file("/usr/local/bin/run_borg_backup@transition-repo")
        prune = host.file("/usr/local/bin/run_borg_prune@transition-repo")

        assert "borg create" in backup.content_string
        assert "borg prune" not in backup.content_string
        assert "borg prune" in prune.content_string
        assert "borg compact" in prune.content_string
        assert "--keep-last 2" in prune.content_string
        assert "--glob-archives '{hostname}-*'" in prune.content_string

    def test_backup_propagates_create_failure_status(self, host):
        if host.backend.get_hostname() != "borg-client":
            return

        result, calls = run_script_with_fake_borg(
            host,
            "/usr/local/bin/run_borg_backup@borg-server",
            "create",
            42,
        )

        assert result.rc == 42
        assert calls == ["create"]

    def test_after_backup_preserves_borg_warning_status(self, host):
        if host.backend.get_hostname() != "borg-client":
            return

        result, calls = run_script_with_fake_borg(
            host,
            "/usr/local/bin/run_borg_backup@borg-server",
            "create",
            1,
        )

        assert result.rc == 1
        assert calls == ["create"]

    @pytest.mark.parametrize(
        ("failed_command", "failed_status", "expected_calls"),
        [
            ("prune", 43, ["prune"]),
            ("compact", 44, ["prune", "compact"]),
        ],
    )
    def test_timer_prune_propagates_failure_status(
        self, host, failed_command, failed_status, expected_calls
    ):
        if host.backend.get_hostname() != "borg-client-2":
            return

        result, calls = run_script_with_fake_borg(
            host,
            "/usr/local/bin/run_borg_prune@borg-server-2",
            failed_command,
            failed_status,
        )

        assert result.rc == failed_status
        assert calls == expected_calls

    def test_backup_script_contains_compression(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert script1.contains("-C zstd")
            assert script2.contains("-C lz4")
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert script1.contains("-C zstd")
            assert script2.contains("-C lz4")
        elif hostname in (
            "borg-client",
            "borg-client-2",
            "borg-client-nonroot",
            "borg-client-transition",
        ):
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.contains("-C")
            assert script.contains("lz4" if hostname == "borg-client-2" else "zstd")
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")

    def test_backup_script_contains_repo_path(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert script1.contains("borg@borg-server")
            assert script1.contains("/opt/borg/configs")
            assert script2.contains("borg@borg-server")
            assert script2.contains("/opt/borg/home-data")
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert script1.contains("borg@borg-server")
            assert script1.contains("/opt/borg/configs-keys")
            assert script2.contains("borg@borg-server")
            assert script2.contains("/opt/borg/home-data-keys")
        elif hostname == "borg-client-nonroot":
            script = host.file("/usr/local/bin/run_borg_backup@borg-server")
            assert script.contains("borg@borg-server")
            assert script.contains("/opt/borg")
        elif hostname == "borg-client":
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.contains("borg@borg-server")
            assert script.contains("/opt/borg")
        elif hostname == "borg-client-2":
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.contains("backupserver@borg-server-2")
            assert script.contains("/var/backups")
        elif hostname == "borg-client-transition":
            script = host.file("/usr/local/bin/run_borg_backup@transition-repo")
            assert script.contains("borg@borg-server")
            assert script.contains("/opt/borg/transition-repo")
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")

    def test_backup_script_contains_backup_paths(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert "/etc" in script1.content_string
            assert "/home" in script2.content_string
            assert "--exclude" in script2.content_string
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert "/etc" in script1.content_string
            assert "/home" in script2.content_string
            assert "--exclude" in script2.content_string
        elif hostname == "borg-client-nonroot":
            script = host.file("/usr/local/bin/run_borg_backup@borg-server")
            content = script.content_string
            assert "/etc" in content
        elif hostname in ("borg-client", "borg-client-2", "borg-client-transition"):
            script = host.file(get_backup_script_paths(hostname)[0])
            content = script.content_string
            assert "/etc" in content or "/home" in content
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")

    def test_backup_script_is_executable(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data")
            assert script1.mode == 0o711
            assert script2.mode == 0o711
        elif hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert script1.mode == 0o711
            assert script2.mode == 0o711
        elif hostname == "borg-client-nonroot":
            script = host.file("/usr/local/bin/run_borg_backup@borg-server")
            assert script.mode == 0o711
        elif hostname in ("borg-client", "borg-client-2", "borg-client-transition"):
            script = host.file(get_backup_script_paths(hostname)[0])
            assert script.mode == 0o711
        else:
            pytest.fail(f"Unexpected hostname: {hostname}")

    def test_backup_script_contains_additional_arguments(self, host):
        hostname = host.backend.get_hostname()

        if hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert "--stats" in script1.content_string
            assert "--list" in script1.content_string
            assert "--filter=AME" in script1.content_string
            assert "--one-file-system" in script2.content_string
            assert "--exclude-caches" in script2.content_string
        elif hostname == "borg-client-2":
            script = host.file(get_backup_script_paths(hostname)[0])
            content = script.content_string
            assert "-C zlib,6" in content

    def test_backup_script_compression_override(self, host):
        hostname = host.backend.get_hostname()
        if hostname != "borg-client-2":
            return

        script = host.file(get_backup_script_paths(hostname)[0])
        content = script.content_string

        assert "-C lz4" in content
        assert "-C zlib,6" in content

    def test_backup_script_contains_borg_rsh_when_per_repo(self, host):
        hostname = host.backend.get_hostname()
        per_repo = get_client_ssh_key_per_repo(host)

        if not per_repo:
            return

        if hostname == "borg-client-multi-keys":
            script1 = host.file("/usr/local/bin/run_borg_backup@configs-keys")
            script2 = host.file("/usr/local/bin/run_borg_backup@home-data-keys")
            assert "BORG_RSH" in script1.content_string
            assert "BORG_RSH" in script2.content_string
            assert "ssh -i" in script1.content_string
            assert "ssh -i" in script2.content_string


class TestAggregateScripts:
    def test_legacy_aggregate_backup_script_is_preserved(self, host):
        aggregate = host.file("/usr/local/bin/run_borg_backup")

        if host.backend.get_hostname() == "borg-client-transition":
            assert aggregate.exists
            assert "legacy-backup-block" in aggregate.content_string
        else:
            assert not aggregate.exists

    def test_aggregate_prune_script_is_not_created(self, host):
        assert not host.file("/usr/local/bin/run_borg_prune").exists