#!/usr/bin/python

import psutil
import argparse
import sys

def main():
    parser = argparse.ArgumentParser(description='check_memory')
    parser.add_argument('-w', '--warning', required=True, type=int,
                        help='Warning threshold percentage (e.g. 70)')
    parser.add_argument('-c', '--critical', required=True, type=int,
                        help='Critical threshold percentage (e.g. 90)')
    args = parser.parse_args()

    # Get the memory usage percentage
    memory = psutil.virtual_memory()
    usage_percent = memory.percent

    # Compare with thresholds and output in Nagios plugin format
    if usage_percent >= args.critical:
        print(f"CRITICAL - Memory usage is {usage_percent}% | memory_usage={usage_percent}%")
        sys.exit(2)
    elif usage_percent >= args.warning:
        print(f"WARNING - Memory usage is {usage_percent}% | memory_usage={usage_percent}%")
        sys.exit(1)
    else:
        print(f"OK - Memory usage is {usage_percent}% | memory_usage={usage_percent}%")
        sys.exit(0)

if __name__ == '__main__':
    main()