############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import pytest

import atf

test_user = atf.properties["test-user"]
slurm_user = atf.properties["slurm-user"]
test_acct = "test_acct"
extra_acct = "extra_acct"
parent_acct = "parent_acct"
restricted_acct = "restricted_acct"


TEST_CASES = [
    dict(
        name="user_match1",
        res_account=None,
        res_user=test_user,
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="user_match2",
        res_account=None,
        res_user=f"{test_user},{slurm_user}",
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="user_match3",
        res_account=None,
        res_user=f"{slurm_user},{test_user}",
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="no_user_match",
        res_account=None,
        res_user=slurm_user,
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="account_match1",
        res_account=test_acct,
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="account_match2",
        res_account=f"{test_acct},{restricted_acct}",
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="account_match3",
        res_account=f"{restricted_acct},{test_acct}",
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="no_account_match",
        res_account=restricted_acct,
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="parent_account_match",
        res_account=parent_acct,
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="deny_user1",
        res_account=None,
        res_user=f"-{test_user}",
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="deny_user2",
        res_account=None,
        res_user=f"-{test_user},-{slurm_user}",
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="deny_user3",
        res_account=None,
        res_user=f"-{slurm_user}",
        expect_invalid_res=False,
        expect_show_res=True,
        expect_job_in_res=True,
    ),
    dict(
        name="acct_match_deny_user",
        res_account=test_acct,
        res_user=f"-{test_user}",
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="deny_acct",
        res_account=f"-{test_acct}",
        res_user=None,
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="user_match_deny_acct",
        res_account=f"-{test_acct}",
        res_user=test_user,
        expect_invalid_res=False,
        expect_show_res=False,
        expect_job_in_res=False,
    ),
    dict(
        name="mixed_access_invalid_res1",
        res_account=None,
        res_user=f"-{test_user},{slurm_user}",
        expect_invalid_res=True,
    ),
    dict(
        name="mixed_access_invalid_res2",
        res_account=None,
        res_user=f"{test_user},-{slurm_user}",
        expect_invalid_res=True,
    ),
    dict(
        name="mixed_access_invalid_res3",
        res_account=f"{test_acct},-{restricted_acct}",
        res_user=None,
        expect_invalid_res=True,
    ),
    dict(
        name="mixed_access_invalid_res4",
        res_account=f"-{test_acct},{restricted_acct}",
        res_user=None,
        expect_invalid_res=True,
    ),
]


@pytest.fixture(scope="module", autouse=True)
def setup():
    if test_user == slurm_user:
        pytest.skip(
            f"This test requires SlurmUser ({slurm_user}) and SlurmTestUser ({test_user}) to be different.",
            allow_module_level=True,
        )

    atf.require_accounting(modify=True)
    atf.require_config_parameter_includes("PrivateData", "reservations")
    atf.require_config_parameter_includes("AccountingStorageEnforce", "associations")
    atf.require_slurm_running()


@pytest.fixture(scope="module", autouse=True)
def setup_db():
    accounts = [
        f"{parent_acct}",
        f"{restricted_acct}",
        f"{test_acct} parent={parent_acct}",
        f"{extra_acct}",
    ]

    users = [
        f"{slurm_user} account={restricted_acct},{extra_acct}",
        f"{test_user} account={test_acct}",
    ]

    try:
        for account in accounts:
            atf.run_command(
                f"sacctmgr -i add account {account}",
                user=slurm_user,
                fatal=True,
            )
        for user in users:
            atf.run_command(
                f"sacctmgr -i add user {user}",
                user=slurm_user,
                fatal=True,
            )

        atf.run_command("sacctmgr show assoc tree", user=slurm_user, fatal=True)
        yield
    finally:
        atf.run_command("sacctmgr show assoc tree", user=slurm_user, fatal=True)
        for user in reversed(users):
            atf.run_command(
                f"sacctmgr -i delete user {user}",
                user=slurm_user,
                fatal=False,
            )
        for account in reversed(accounts):
            atf.run_command(
                f"sacctmgr -i delete account {account}",
                user=slurm_user,
                fatal=False,
            )


@pytest.fixture
def reservation(request):
    try:
        case = request.param
        res_name = f"res_{case['name']}"

        # Create reservation
        res_args = [
            f"reservationname={res_name}",
            "start=now",
            "duration=10",
            "nodecnt=1",
        ]
        if case.get("res_user"):
            res_args.append(f"user={case['res_user']}")
        if case.get("res_account"):
            res_args.append(f"account={case['res_account']}")

        create_cmd = f"scontrol create {' '.join(res_args)}"
        result = atf.run_command(create_cmd, user=slurm_user)
        if case["expect_invalid_res"]:
            assert (
                result["exit_code"] != 0
            ), f"Should not be able to create invalid reservation {res_name}"
            res_name = None
        else:
            assert result["exit_code"] == 0, f"Couldn't create reservation {res_name}"

        # Yield values for the test
        yield case, res_name

    finally:
        atf.cancel_all_jobs(fatal=True)

        if res_name:
            atf.run_command(
                f"scontrol delete reservationname={res_name}", user=slurm_user
            )


@pytest.fixture
def extra_assoc():
    atf.run_command(
        f"sacctmgr -i add user {test_user} account={extra_acct}",
        user=slurm_user,
        fatal=True,
    )
    yield
    atf.run_command(
        f"sacctmgr -i delete user {test_user} account={extra_acct}",
        user=slurm_user,
        fatal=False,
    )


@pytest.fixture
def show_submit_resv():
    resv = "res_show_submit"
    atf.run_command(
        f"scontrol create reservationname={resv} start=now duration=10 "
        f"nodecnt=1 account=-{test_acct} user={test_user}",
        user=slurm_user,
        fatal=True,
    )
    yield resv
    atf.run_command(f"scontrol delete reservationname={resv}", user=slurm_user)


@pytest.fixture
def delete_submit_resv():
    resv = "res_del_submit"
    atf.run_command(
        f"scontrol create reservationname={resv} start=now duration=10 "
        f"nodecnt=1 account=-{test_acct} user={test_user} flags=USER_DELETE",
        user=slurm_user,
        fatal=True,
    )
    yield resv
    # The test deletes the reservation, so it may already be gone at teardown
    atf.run_command(
        f"scontrol delete reservationname={resv}", user=slurm_user, xfail=True
    )


@pytest.fixture
def skip_submit_resv():
    resv = "res_skip_submit"
    # Skipping requires a recurring reservation (DAILY)
    atf.run_command(
        f"scontrol create reservationname={resv} start=now duration=10 "
        f"nodecnt=1 account=-{test_acct} user={test_user} flags=USER_DELETE,DAILY",
        user=slurm_user,
        fatal=True,
    )
    yield resv
    atf.run_command(f"scontrol delete reservationname={resv}", user=slurm_user)


def show_as_slurm_user(res_name):
    show = atf.run_command(f"scontrol show reservation {res_name}", user=slurm_user)
    assert (
        f"ReservationName={res_name} StartTime=" in show["stdout"]
    ), "Reservation should always be visible as SlurmUser"


def show_as_test_user(case, res_name):
    show = atf.run_command(f"scontrol show reservation {res_name}", user=test_user)
    if case["expect_show_res"]:
        assert (
            f"ReservationName={res_name} StartTime=" in show["stdout"]
        ), "Reservation should be visible"
    else:
        assert (
            f"ReservationName={res_name} StartTime=" not in show["stdout"]
        ), "Reservation should not be visible"


def run_job_as_slurm_user(res_name):
    run = atf.run_command(
        f"srun -N1 --account={restricted_acct} --reservation={res_name} true",
        user=slurm_user,
    )
    assert run["exit_code"] == 0, "Job should have been allowed in reservation"


def run_job_as_test_user(case, res_name):
    run = atf.run_command(
        f"srun -N1  --account={test_acct} --reservation={res_name} true", user=test_user
    )
    if case["expect_job_in_res"]:
        assert run["exit_code"] == 0, "Job should have been allowed in reservation"
    else:
        assert run["exit_code"] != 0, "Job should have been denied for reservation"
        assert "Access denied" in run["stderr"]


def assert_submit_allowed(resv, acct):
    assert (
        atf.run_job_exit(
            f"-N1 --account={acct} --reservation={resv} true",
            user=test_user,
        )
        == 0
    ), "Reservation submission with allowed account should succeed"


@pytest.mark.parametrize(
    "reservation", TEST_CASES, indirect=True, ids=lambda c: c["name"]
)
def test_reservation_access(reservation):
    case, res_name = reservation

    if res_name is None:
        return

    show_as_slurm_user(res_name)
    show_as_test_user(case, res_name)
    run_job_as_slurm_user(res_name)
    run_job_as_test_user(case, res_name)


@pytest.mark.xfail(
    atf.get_version("sbin/slurmctld") < (26, 11),
    reason="Ticket 24457: mixed denied/allowed reservation association access fixed in 26.11",
)
def test_show_matches_submit(extra_assoc, show_submit_resv):
    """A user who can submit to a mixed-access reservation can view it."""

    resv = show_submit_resv

    # The denied account cannot submit
    denied = atf.run_job(
        f"-N1 --account={test_acct} --reservation={resv} true",
        user=test_user,
        xfail=True,
    )
    assert (
        denied["exit_code"] != 0
    ), "Reservation submission with denied account should fail"
    assert (
        "Access denied" in denied["stderr"]
    ), f"Expected 'Access denied', got: {denied['stderr']}"

    # The allowed account can submit, so the reservation should be visible
    assert_submit_allowed(resv, extra_acct)
    assert resv in atf.get_reservations(
        user=test_user
    ), "Reservation visibility should match submission access"


@pytest.mark.xfail(
    atf.get_version("sbin/slurmctld") < (26, 11),
    reason="Ticket 24457: mixed denied/allowed reservation association access fixed in 26.11",
)
def test_delete_matches_submit(extra_assoc, delete_submit_resv):
    """A user who can submit to a mixed-access reservation can delete it."""

    resv = delete_submit_resv

    # The allowed account can submit, so the user can delete the reservation
    assert_submit_allowed(resv, extra_acct)
    atf.run_command(
        f"scontrol delete reservationname={resv}",
        user=test_user,
        fatal=True,
    )


@pytest.mark.xfail(
    atf.get_version("sbin/slurmctld") < (26, 11),
    reason="Ticket 24457: mixed denied/allowed reservation association access fixed in 26.11",
)
def test_skip_matches_submit(extra_assoc, skip_submit_resv):
    """A user who can submit to a mixed-access reservation can skip it."""

    resv = skip_submit_resv

    # The allowed account can submit, so the user can skip the reservation
    assert_submit_allowed(resv, extra_acct)
    atf.run_command(
        f"scontrol update reservationname={resv} skip",
        user=test_user,
        fatal=True,
    )