import os
import stat
import subprocess
from pathlib import Path

SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "mem-guard-root.sh"

OOMCTL_DISARMED = """\
Swap Monitored CGroups:
Memory Pressure Monitored CGroups:
\tPath: /user.slice/user-0.slice/user@0.service
"""

OOMCTL_ARMED = """\
Swap Monitored CGroups:
\tPath: /
Memory Pressure Monitored CGroups:
\tPath: /user.slice/user-0.slice/user@0.service
"""


def _write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text)


OOMD_DEFAULTS = "[OOM]\nDefaultMemoryPressureLimit=40%\nDefaultMemoryPressureDurationSec=10s\nSwapUsedLimit=45%\n"


def _sysroot(tmp_path: Path, *, oomctl: str, root_slice: str, memcap: str, swappiness_files: dict) -> Path:
    root = tmp_path / "sysroot"
    _write(root / "etc/systemd/oomd.conf.d/99-early.conf", OOMD_DEFAULTS)
    _write(root / "etc/systemd/system/-.slice.d/10-oomd.conf", root_slice)
    _write(root / "etc/systemd/system/user-.slice.d/10-memcap.conf", memcap)
    _write(root / "proc/sys/vm/swappiness", "10\n")
    for name, text in swappiness_files.items():
        _write(root / "etc/sysctl.d" / name, text)

    fake = tmp_path / "bin" / "oomctl"
    _write(fake, "#!/bin/sh\ncat <<'EOF'\n" + oomctl + "EOF\n")
    fake.chmod(fake.stat().st_mode | stat.S_IEXEC)
    return root


def _audit(tmp_path: Path, root: Path):
    env = dict(os.environ, SYSROOT=str(root), OOMCTL=str(tmp_path / "bin" / "oomctl"))
    return subprocess.run(
        ["bash", str(SCRIPT), "--audit"], env=env, capture_output=True, text=True
    )


def test_clean_posture_passes(tmp_path):
    root = _sysroot(
        tmp_path,
        oomctl=OOMCTL_DISARMED,
        root_slice="[Slice]\nManagedOOMSwap=auto\n",
        memcap="[Slice]\n",
        swappiness_files={"99-swappiness.conf": "vm.swappiness=10\n"},
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 0, res.stdout
    assert "DRIFT" not in res.stdout


def test_swap_kill_armed_is_drift(tmp_path):
    root = _sysroot(
        tmp_path,
        oomctl=OOMCTL_ARMED,
        root_slice="[Slice]\nManagedOOMSwap=kill\n",
        memcap="[Slice]\n",
        swappiness_files={"99-swappiness.conf": "vm.swappiness=10\n"},
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT oomd-swap-kill" in res.stdout
    assert "DRIFT oomd-dropins" in res.stdout


def test_session_wide_cap_is_drift(tmp_path):
    root = _sysroot(
        tmp_path,
        oomctl=OOMCTL_DISARMED,
        root_slice="[Slice]\nManagedOOMSwap=auto\n",
        memcap="[Slice]\nMemoryHigh=48G\nMemoryMax=58G\n",
        swappiness_files={"99-swappiness.conf": "vm.swappiness=10\n"},
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT user-slice-cap" in res.stdout


def test_conflicting_swappiness_files_are_drift(tmp_path):
    root = _sysroot(
        tmp_path,
        oomctl=OOMCTL_DISARMED,
        root_slice="[Slice]\nManagedOOMSwap=auto\n",
        memcap="[Slice]\n",
        swappiness_files={
            "99-swappiness.conf": "vm.swappiness=10\n",
            "99-mem-guard.conf": "vm.swappiness = 150\n",
        },
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT swappiness" in res.stdout


def _agent_conf(root: Path, rel: str, text: str) -> None:
    """agent.slice config as one of its three writers would leave it.

    rel is relative to /etc for root-owned drop-ins, or to $HOME otherwise.
    """
    base = root / "etc" if rel.startswith("systemd/user") else root / str(Path.home()).lstrip("/")
    _write(base / rel, text)


def _clean_sysroot(tmp_path: Path) -> Path:
    return _sysroot(
        tmp_path,
        oomctl=OOMCTL_DISARMED,
        root_slice="[Slice]\nManagedOOMSwap=auto\n",
        memcap="[Slice]\n",
        swappiness_files={"99-swappiness.conf": "vm.swappiness=10\n"},
    )


def test_pressure_only_agent_slice_passes(tmp_path):
    root = _clean_sysroot(tmp_path)
    _agent_conf(
        root,
        "systemd/user/agent.slice.d/90-ceiling.conf",
        "[Slice]\nManagedOOMSwap=auto\nManagedOOMMemoryPressure=kill\nManagedOOMMemoryPressureLimit=60%\n",
    )
    _agent_conf(root, ".config/systemd/user/agent.slice", "[Slice]\nCPUQuota=900%\nTasksMax=3072\n")
    res = _audit(tmp_path, root)
    assert res.returncode == 0, res.stdout
    assert "OK    agent-slice-cap" in res.stdout


def test_root_dropin_memory_ceiling_is_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _agent_conf(root, "systemd/user/agent.slice.d/90-ceiling.conf", "[Slice]\nMemoryMax=52%\n")
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT agent-slice-cap" in res.stdout


def test_set_property_runtime_dropin_is_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _agent_conf(
        root,
        ".config/systemd/user.control/agent.slice.d/50-MemoryHigh.conf",
        "# created via systemctl set-property\n[Slice]\nMemoryHigh=25769803776\n",
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT agent-slice-cap" in res.stdout


def test_agent_slice_swap_kill_is_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _agent_conf(root, "systemd/user/agent.slice.d/90-ceiling.conf", "[Slice]\nManagedOOMSwap=kill\n")
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT agent-slice-cap" in res.stdout


def test_explicitly_unset_ceiling_is_not_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _agent_conf(
        root,
        ".config/systemd/user.control/agent.slice.d/50-MemoryMax.conf",
        "[Slice]\nMemoryMax=infinity\n",
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 0, res.stdout


def test_runtime_user_control_dropin_is_drift(tmp_path):
    """`systemctl --user set-property --runtime` writes under /run, not $HOME."""
    root = _clean_sysroot(tmp_path)
    _write(
        root / f"run/user/{os.getuid()}/systemd/user.control/agent.slice.d/50-MemoryMax.conf",
        "[Slice]\nMemoryMax=34359738368\n",
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT agent-slice-cap" in res.stdout


def test_edited_oomd_defaults_are_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _write(
        root / "etc/systemd/oomd.conf.d/99-early.conf",
        OOMD_DEFAULTS.replace("40%", "60%"),
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT oomd-defaults" in res.stdout


def test_missing_oomd_defaults_are_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    (root / "etc/systemd/oomd.conf.d/99-early.conf").unlink()
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT oomd-defaults" in res.stdout


def _owner_cgroup(root: Path, *slices: str) -> Path:
    uid = os.getuid()
    base = root / f"sys/fs/cgroup/user.slice/user-{uid}.slice/user@{uid}.service"
    for s in slices:
        (base / s).mkdir(parents=True, exist_ok=True)
    return base


def test_owner_cgroup_without_omit_is_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _owner_cgroup(root, "app.slice", "human.slice")
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT owner-oom-omit" in res.stdout
    assert "app.slice" in res.stdout and "human.slice" in res.stdout


def test_owner_cgroup_with_omit_xattr_passes(tmp_path):
    root = _clean_sysroot(tmp_path)
    base = _owner_cgroup(root, "app.slice", "human.slice")
    for s in ("app.slice", "human.slice"):
        try:
            os.setxattr(base / s, "user.oomd_omit", b"1")
        except OSError as exc:
            import pytest

            pytest.skip(f"user xattrs unsupported on {tmp_path}: {exc}")
    res = _audit(tmp_path, root)
    assert res.returncode == 0, res.stdout
    assert "OK    owner-oom-omit" in res.stdout


def test_user_manager_swap_kill_dropin_is_drift(tmp_path):
    root = _clean_sysroot(tmp_path)
    _write(
        root / "etc/systemd/system/user@.service.d/50-oomd.conf",
        "[Service]\nManagedOOMSwap=kill\nManagedOOMMemoryPressure=kill\n",
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 1
    assert "DRIFT oomd-dropins" in res.stdout


def test_user_manager_pressure_only_dropin_passes(tmp_path):
    root = _clean_sysroot(tmp_path)
    _write(
        root / "etc/systemd/system/user@.service.d/50-oomd.conf",
        "[Service]\nManagedOOMSwap=auto\nManagedOOMMemoryPressure=kill\n",
    )
    res = _audit(tmp_path, root)
    assert res.returncode == 0, res.stdout
