#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import json
import memgraph_plugins
import mem_utils
import getopt

if os.geteuid() != 0:
    print("This program must be run as root. Aborting.")
    sys.exit(1)


def memgraph_handler_cmd(argv):
    try:
        opts, args = getopt.getopt(argv, "j:hgfsakli:c:p:d")
    except getopt.GetoptError:
        print("get opt error")
        sys.exit(2)

    enable_plugins = [
        {"name": "memcgLeak", "params": {}},
        {"name": "memfrag", "params": {}},
    ]

    json_path = ""
    enable_memleak = False
    enable_pod = False
    memleak_detect_duration = 60
    for opt, arg in opts:
        if opt == "-h":
            print("-g: only show memory usage graph")
            print("-f: show file cache usage detail")
            print("-k: check kernel memleak")
            print("-i <duration>: set memleak detect duration")
            print("-a: show anon page usage detail")
            print("-l: show memory usage of thread")
            # print("-c: show memory usage of cgroup")
            print("-j: output for json file")
            print("-p: memgraph for specific pod")
            # print("-d: show memory usage graph and diagnosis")
            sys.exit()
        elif opt in ("-g"):
            # 启用 memgraph 插件，展示内存使用情况
            enable_plugins.append({"name": "memgraph", "params": {}})
        elif opt in ("-j"):
            json_path = arg
        elif opt in ("-f"):
            # 启用 cache 插件，扫描系统文件缓存已经共享内存文件的占用情况
            enable_plugins.append({"name": "cache", "params": {}})
        elif opt in ("-a"):
            # 启用 anonTaskTop 插件，展示匿名内存使用最多的前 30 个进程
            enable_plugins.append({"name": "anonTaskTop", "params": {"topN": 30}})
        elif opt in ("-l"):
            # 启用 vmRssTaskTop 插件，展示 VmRSS 内存使用最多的前 30 个进程
            enable_plugins.append({"name": "vmRssTaskTop", "params": {"topN": 30}})
        elif opt in ("-k"):
            enable_memleak = True
        elif opt in ("-i"):
            memleak_detect_duration = int(arg)
            if memleak_detect_duration < 0:
                memleak_detect_duration = 60
            elif memleak_detect_duration > 300:
                memleak_detect_duration = 300
        elif opt in ("-c"):
            # Ignore
            pass
        elif opt in ("-d"):
            # 启用 abnormalPackages，扫描系统上是否安装了可能存在内存泄露的包
            enable_plugins.append({"name": "abnormalPackages", "params": {}})
            pass
        elif opt in ("-p"):
            pod_name = arg
            # if enable pod mode, make sure podDetect plugin should be executed first
            enable_plugins.insert(0, {"name": "podDetect", "params": {"pod_name": pod_name}})
            mem_utils.set_podinfo_lib()
            enable_pod = True

    if enable_memleak:
        enable_plugins.append(
            {
                "name": "memleak",
                "params": {"memleak_detect_duration": memleak_detect_duration},
            }
        )

    res = memgraph_plugins.execute_plugins(enable_plugins, enable_pod)
    if not json_path:
        print(res.generate_render_text())
    else:
        with open(json_path, "w") as f:
            f.write(json.dumps(res.generate_json_result()))
    return res


if __name__ == "__main__":
    memgraph_handler_cmd(sys.argv[1:])
