#!/usr/bin/python
import argparse
import ast
import sys
import syslog

import psutil
from sonic_py_common import multi_asic
import swsssdk


def check_process_existence(container_name, process_cmdline):
    """
    @summary: Check whether the process in the specified container is running or not and
              an alerting message will written into syslog if it failed to run.
    """
    config_db = swsssdk.ConfigDBConnector()
    config_db.connect()
    feature_table = config_db.get_table("FEATURE")

    if container_name in feature_table.keys():
        # We look into the 'FEATURE' table to verify whether the container is disabled or not.
        # If the container is diabled, we exit.
        if ("state" in feature_table[container_name].keys()
                and feature_table[container_name]["state"] == "disabled"):
            sys.exit(0)
        else:
            # We leveraged the psutil library to help us check whether the process is running or not.
            # If the process entity is found in process tree and it is also in the 'running' or 'sleeping'
            # state, then it will be marked as 'running'.

            # For given feature we get the host and network namespace instances it's processes should be running
            # based on it's scope and add it to expected set.

            # From psutil we get number of running instances of the processes and add it to the the actual set

            # Difference bwetween expected and actual set provides instances where the processes are not running
            # and will be logged as syslog message by monit

            process_namespace_expected_set = set()
            process_namespace_found_set = set()

            has_global_scope = ast.literal_eval(feature_table[container_name].get('has_global_scope', 'True'))
            has_per_asic_scope = ast.literal_eval(feature_table[container_name].get('has_per_asic_scope', 'False'))

            if has_global_scope:
                process_namespace_expected_set.add(multi_asic.DEFAULT_NAMESPACE)

            if has_per_asic_scope:
                process_namespace_expected_set.update(multi_asic.get_namespace_list())

            for process in psutil.process_iter(["cmdline", "status", "pid"]):
                try:
                    if ((' '.join(process.cmdline())).startswith(process_cmdline) and process.status() in ["running", "sleeping"]):
                         process_namespace_found_set.add(multi_asic.get_current_namespace(process.info['pid']))
                except psutil.NoSuchProcess:
                    pass

            process_namespace_diff_set = process_namespace_expected_set.difference(process_namespace_found_set)

            if process_namespace_diff_set:
                 host_display_str = ""
                 namespace_display_str = ""

                 for ns in process_namespace_diff_set:
                     if ns == multi_asic.DEFAULT_NAMESPACE:
                         host_display_str = " in host"
                     else:
                         if not namespace_display_str:
                             namespace_display_str = " in namespace " + ns
                         else:
                             namespace_display_str += ", " + ns

                 join_str = " and" if host_display_str and namespace_display_str else ""
                 
                 # If this script is run by Monit, then the following output will be appended to
                 # Monit's syslog message.
                 print("'{}' is not running{}{}{}".format(process_cmdline, host_display_str, join_str, namespace_display_str))
                 sys.exit(1)
    else:
        syslog.syslog(syslog.LOG_ERR, "container '{}' is not included in SONiC image or the given container name is invalid!"
                      .format(container_name))


def main():
    parser = argparse.ArgumentParser(description="Check whether the process in the specified \
             container is running and an alerting message will be written into syslog if it \
             failed to run.", usage="/usr/bin/process_checker <container_name> <process_cmdline>")
    parser.add_argument("container_name", help="container name")
    parser.add_argument("process_cmdline", nargs=argparse.REMAINDER, help="process command line")
    args = parser.parse_args()

    check_process_existence(args.container_name, ' '.join(args.process_cmdline))


if __name__ == '__main__':
    main()