#!/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 sys

import lib.base
import lib.human
import lib.version
from lib.globals import (STATE_OK, STATE_UNKNOWN)

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


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

DESCRIPTION = """Displays system memory usage and alerts on sustained high usage.
Reports total/used/available/free plus shared/buffers/cached, and evaluates WARN/CRIT
against the overall usage percentage. Perfdata is emitted for all fields so you can
graph trends over time. With `--top`, the most memory-consuming processes are listed
(by RSS and percentage) to aid quick diagnosis. Cross-platform on all psutil-supported
systems (Linux, Windows, *BSD, macOS)."""

DEFAULT_CRIT = 95  # %
DEFAULT_TOP = 5
DEFAULT_WARN = 90  # %


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='Always returns OK.',
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

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

    parser.add_argument(
        '--top',
        help='List x "Top most memory consuming processes". '
             'Use `--top=0` to disable this feature. '
             'Default: %(default)s',
        dest='TOP',
        type=int,
        default=DEFAULT_TOP,
    )

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

    return parser.parse_args()


def top(count):
    """Get top X most memory consuming processes.
    """
    # Fast path: nothing to print, so nothing to scan
    if count <= 0:
        return ''

    procs = {}  # name > {'%': float, 'rss': int}
    msg = ''

    # Prefer attrs path (psutil >= 5.3.0): fewer syscalls, fewer exceptions
    if lib.version.version(psutil.__version__) >= lib.version.version('5.3.0'):
        try:
            for p in psutil.process_iter(attrs=['name', 'memory_percent', 'memory_info'], ad_value=None):  # pylint: disable=C0301
                try:
                    info = p.info
                    name = info.get('name') or ''
                    # On some platforms memory_percent can be None briefly; guard it.
                    mem_pct = float(info.get('memory_percent') or 0.0)
                    mi = info.get('memory_info')
                    rss = getattr(mi, 'rss', 0) if mi is not None else 0
                    entry = procs.setdefault(name, {'%': 0.0, 'rss': 0})
                    entry['%'] += mem_pct
                    entry['rss'] += rss
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    continue
        except Exception:
            # Defensive: if psutil attrs path misbehaves on some platform/version, fall back below.
            pass

    # Legacy / fallback path
    if not procs:
        try:
            for proc in psutil.process_iter():
                try:
                    info = proc.as_dict(attrs=['name', 'memory_percent', 'memory_info'])
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    continue
                name = info.get('name') or ''
                mem_pct = float(info.get('memory_percent') or 0.0)
                mi = info.get('memory_info')
                rss = getattr(mi, 'rss', 0) if mi is not None else 0
                entry = procs.setdefault(name, {'%': 0.0, 'rss': 0})
                entry['%'] += mem_pct
                entry['rss'] += rss
        except psutil.NoSuchProcess:
            pass

    if not procs:
        return msg

    # Sort by percentage, desc; produce up to 'count'
    ranked = sorted(procs.items(), key=lambda kv: kv[1]['%'], reverse=True)[:count]
    lines = [f'\nTop {count} most memory consuming processes:']
    for i, (name, agg) in enumerate(ranked, start=1):
        lines.append(f'{i}. {name}: {lib.human.bytes2human(agg["rss"])} ({agg["%"]:.1f}%)')
    return '\n'.join(lines) + '\n'


def main():
    """The main function. Hier spielt die Musik.
    """

    # parse the command line, exit with UNKNOWN if it fails
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    # fetch data
    virt = psutil.virtual_memory()

    # init some vars
    state = STATE_OK
    perfdata = ''

    # analyze data and build the message
    usage_percent = float(getattr(virt, 'percent', 0))
    state = lib.base.get_state(usage_percent, args.WARN, args.CRIT, _operator='ge')

    # Header: percentage + main sizes
    header_parts = [
        f'{usage_percent:.0f}%{lib.base.state2str(state, prefix=" ")}',
        f'total: {lib.human.bytes2human(getattr(virt, "total", 0))}',
        f'used: {lib.human.bytes2human(getattr(virt, "used", 0))}',
        f'available: {lib.human.bytes2human(getattr(virt, "available", 0))}',
        f'free: {lib.human.bytes2human(getattr(virt, "free", 0))}',
    ]
    msg_header = ' - '.join([header_parts[0], ', '.join(header_parts[1:])])

    # Body: extended fields (if present)
    body_parts = [
        f'shared: {lib.human.bytes2human(getattr(virt, "shared", 0))}',
        f'buffers: {lib.human.bytes2human(getattr(virt, "buffers", 0))}',
        f'cached: {lib.human.bytes2human(getattr(virt, "cached", 0))}',
    ]
    msg_body = '\n' + ', '.join(body_parts)

    # Top X most memory consuming processes
    msg_body += '\n' + top(args.TOP)

    # perfdata (identical keys/units/min/max as before)
    total = getattr(virt, 'total', 0)
    perfdata += lib.base.get_perfdata(
        'usage_percent',
        usage_percent,
        uom='%',
        warn=args.WARN,
        crit=args.CRIT,
        _min=0,
        _max=100,
    )
    perfdata += lib.base.get_perfdata(
        'total',
        total,
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'used',
        getattr(virt, 'used', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'available',
        getattr(virt, 'available', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'free',
        getattr(virt, 'free', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'shared',
        getattr(virt, 'shared', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'buffers',
        getattr(virt, 'buffers', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )
    perfdata += lib.base.get_perfdata(
        'cached',
        getattr(virt, 'cached', 0),
        uom='B',
        warn=None,
        crit=None,
        _min=0,
        _max=total,
    )

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


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