#!/usr/bin/env python3
# -*- coding: utf-8; py-indent-offset: 4 -*-
#
# Author:  Linuxfabrik GmbH, Zurich, Switzerland
# Contact: info (at) linuxfabrik (dot) ch
#          https://www.linuxfabrik.ch/
# License: The Unlicense, see LICENSE file.

# https://github.com/Linuxfabrik/monitoring-plugins/blob/main/CONTRIBUTING.md

"""See the check's README for more details."""

import argparse
import json
import os
import sys

import lib.args
import lib.base
import lib.human
import lib.lftest
import lib.txt
from lib.globals import STATE_UNKNOWN

try:
    import psutil
except ImportError:
    print('Python module "psutil" is not installed.')
    sys.exit(STATE_UNKNOWN)


__author__ = 'Linuxfabrik GmbH, Zurich/Switzerland'
__version__ = '2026052001'

DESCRIPTION = """Monitors swap space usage as a percentage of total swap. Optionally lists the top
processes consuming the most swap to help identify the source of high usage. Alerts
when usage exceeds the configured thresholds. Returns a configurable severity when no
swap is configured at all, which helps detect a swap partition that was inadvertently
disabled."""

DEFAULT_CRIT = 90  # %
DEFAULT_NO_SWAP = 'unknown'
DEFAULT_TOP = 5
DEFAULT_WARN = 70  # %


def parse_args():
    """Parse command line arguments using argparse."""
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    parser.add_argument(
        '-V',
        '--version',
        action='version',
        version=f'%(prog)s: v{__version__} by {__author__}',
    )

    parser.add_argument(
        '--always-ok',
        help=lib.args.help('--always-ok'),
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c',
        '--critical',
        default=DEFAULT_CRIT,
        dest='CRIT',
        help='Threshold for swap usage, in percent. Default: %(default)s',
        type=int,
    )

    parser.add_argument(
        '--severity-no-swap',
        help='Severity for alerting if no swap is configured. Default: %(default)s',
        dest='SEVERITY_NO_SWAP',
        default=DEFAULT_NO_SWAP,
        choices=['ok', 'warn', 'crit', 'unknown'],
    )

    parser.add_argument(
        '--top',
        help='Number of top processes consuming the most swap space to list (not available on Windows). '
        'Default: %(default)s',
        dest='TOP',
        type=int,
        default=DEFAULT_TOP,
    )

    parser.add_argument(
        '--test',
        help=lib.args.help('--test'),
        dest='TEST',
        type=lib.args.csv,
    )

    parser.add_argument(
        '-w',
        '--warning',
        default=DEFAULT_WARN,
        dest='WARN',
        help='Threshold for swap usage, in percent. Default: %(default)s',
        type=int,
    )

    args, _ = parser.parse_known_args()
    return args


def _load_swap_fixture(raw_json):
    """Convert a test fixture into the shape the plugin expects from
    `psutil.swap_memory()`. The plugin accesses fields as attributes
    (`swap.percent`, `swap.total`, ...) so we build a lightweight
    object with the required attributes. Missing fields default to 0.

    Fixture shape:

        {"total": <bytes>, "used": <bytes>, "free": <bytes>,
         "percent": <0..100>, "sin": <bytes>, "sout": <bytes>}
    """

    class _Swap:
        pass

    data = json.loads(raw_json)
    swap = _Swap()
    for key in ('total', 'used', 'free', 'percent', 'sin', 'sout'):
        setattr(swap, key, data.get(key, 0))
    return swap


def get_swap_and_name(pid):
    """Scan for usage details in /proc-files."""
    try:
        with open(f'/proc/{pid}/status') as file:
            content = file.read()
            name = lib.txt.extract_str(content, 'Name:', '\n')
            size = lib.txt.extract_str(content, 'VmSwap:', '\n')
            if not size:
                size = '0 kB'
            return name.strip(), int(size.replace(' kB', '').strip())
    except FileNotFoundError:
        return None, 0
    except ProcessLookupError:
        return None, 0


def top(count):
    """Get top X processes causing swap usage."""
    process_info = {}
    for pid in os.listdir('/proc'):
        if pid.isdigit():
            name, usage = get_swap_and_name(pid)  # usage is in kB
            try:
                process_info[name] += usage * 1024
            except Exception:
                process_info[name] = usage * 1024
    msg = f'\nTop {count} processes that use the most swap space:\n'
    process_info = lib.base.sort(process_info)
    for i, p in enumerate(process_info[:count]):
        msg += f'{i + 1}. {p[0]}: {lib.human.bytes2human(p[1])}\n'
    return msg


def main():
    """The main function. This is where the magic happens."""

    # parse the command line
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    # fetch data
    if args.TEST is None:
        try:
            swap = psutil.swap_memory()
        except RuntimeError:
            lib.base.oao(
                'Performance counters may be corrupt or disabled. If the counters are corrupt, attempt to [rebuild them](https://docs.microsoft.com/en-us/troubleshoot/windows-server/performance/rebuild-performance-counter-library-values). If the counters are disabled, check the registry `HKLM\\SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Perflib` and make sure there is no key "Disable" (or set it to 0).',
                STATE_UNKNOWN,
            )
    else:
        stdout, _, _ = lib.lftest.test(args.TEST)
        swap = _load_swap_fixture(stdout)
    swap_usage_percent = swap.percent

    # build the message
    msg_header = (
        f'{swap_usage_percent}% - '
        f'total: {lib.human.bytes2human(swap.total)}, '
        f'used: {lib.human.bytes2human(swap.used)}, '
        f'free: {lib.human.bytes2human(swap.free)}'
    )

    perfdata = lib.base.get_perfdata(
        'usage_percent',
        swap_usage_percent,
        uom='%',
        warn=args.WARN,
        crit=args.CRIT,
        _min=0,
        _max=100,
    )
    perfdata += lib.base.get_perfdata(
        'total',
        swap.total,
        uom='B',
        _min=0,
        _max=swap.total,
    )
    perfdata += lib.base.get_perfdata(
        'used',
        swap.used,
        uom='B',
        _min=0,
        _max=swap.total,
    )
    perfdata += lib.base.get_perfdata(
        'free',
        swap.free,
        uom='B',
        _min=0,
        _max=swap.total,
    )

    msg_body = ''
    if not lib.base.WINDOWS:
        msg_body = (
            f'swapped in: '
            f'{lib.human.bytes2human(getattr(swap, "sin", 0))}'
            f', swapped out: '
            f'{lib.human.bytes2human(getattr(swap, "sout", 0))}'
            f' (both cumulative)'
        )
        perfdata += lib.base.get_perfdata(
            'sin',
            getattr(swap, 'sin', 0),
            uom='B',
            _min=0,
        )
        perfdata += lib.base.get_perfdata(
            'sout',
            getattr(swap, 'sout', 0),
            uom='B',
            _min=0,
        )

    # On Linux only, get the top processes causing swap usage
    # (skip under --test to keep the test deterministic and avoid
    # walking the real /proc tree)
    if args.TEST is None and not lib.base.WINDOWS and swap_usage_percent > 0:
        msg_body += '\n' + top(args.TOP)

    # no swap configured at all: report with the configured severity so an
    # inadvertently disabled swap partition does not silently pass as OK
    if swap.total == 0:
        lib.base.oao(
            'No swap configured.',
            lib.base.str2state(args.SEVERITY_NO_SWAP),
            perfdata,
            always_ok=args.ALWAYS_OK,
        )

    # calculating the final check state
    state = lib.base.get_state(swap_usage_percent, args.WARN, args.CRIT)

    # over and out
    lib.base.oao(
        msg_header + '\n' + msg_body, state, perfdata, always_ok=args.ALWAYS_OK
    )


if __name__ == '__main__':
    try:
        main()
    except Exception:
        lib.base.cu()
