#! /usr/bin/python3.7
# Copyright © 2021 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation
# and its affiliates (the "Company") and all right, title, and interest
# in and to the software product, including all associated intellectual
# property rights, are and shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
# All Rights reserved.

import daemon
from core.infra.cue import config, main
import argparse
import logging
import logging.handlers
import os
import signal
import faulthandler


# TODO: NOTE: See note at the end of the file for further signal
# handling related discussion.
# On SIGSEGV, SIGFPE, SIGABRT, SIGBUS, and SIGILL signals, dump stack.  A bit
# more useful than just a core dump.
faulthandler.enable()

# Also, dump core on a SIGUSR1.  This is useful when debugging a deadlock.
faulthandler.register(signal.SIGUSR1)


# Parse our arguments
parser = argparse.ArgumentParser(description="cue: management daemon")
parser.add_argument("-d", "--daemon",
                    help="Run as a daemon",
                    action="store_true")
parser.add_argument("--debug",
                    help="Log DEBUG messages",
                    action="store_true")
parser.add_argument("--logfile",
                    help=(
                      "Write logs to the given file.  If '-', logs are "
                      "written to stdout.  By default, logs are "
                      "written to syslog when in daemon mode and stdout "
                      "otherwise."
                    ),
                    action="store")
parser.add_argument("--profile",
                    help=(
                      "Enables detailed profiling of nvued. This dumps "
                      "profiling data to /tmp/nvued_profile. This has "
                      "performance impact and should only be used for "
                      "debugging purposes. "
                    ),
                    action="store_true")
parser.add_argument("--pruning",
                    help=(
                      "Enables dynamic pruning of the OM schema"
                      "based on the pruning tags supported on the"
                      "system."
                    ),
                    action="store_true")

args = parser.parse_args()


def logger_init_syslog(loglevel):
    syslog_h = logging.handlers.SysLogHandler(address="/dev/log")
    formatter = logging.Formatter("nvued: %(levelname)7s: %(filename)s:%(lineno)s %(message)s")
    syslog_h.setFormatter(formatter)
    syslog_h.setLevel(loglevel)

    # Add a filter to remove chatty systemd-notify messages
    if 'WATCHDOG_PID' in os.environ:
        syslog_h.addFilter(lambda r: 0 if 'systemd-notify' in r.msg else 1)
    logging.getLogger().addHandler(syslog_h)


def logger_init_file(loglevel, logfile):
    # If logfile is "-", write to stdout.  This is a lot like
    # logger_init_stdout(), but we don't try to color the output.  This option
    # is useful when piping into log rotators.
    if logfile == "-":
        logging.basicConfig(
            level=loglevel,
            format="%(asctime)s %(levelname)7s: %(name)s: %(filename)s:%(lineno)s %(message)s")
    else:
        logging.basicConfig(
            filename=logfile,
            filemode='a',
            level=loglevel,
            format="%(asctime)s %(levelname)7s: %(name)s: %(filename)s:%(lineno)s %(message)s")


def logger_init_stdout(loglevel, logfile=None):
    logging.basicConfig(
        level=loglevel,
        format="%(asctime)s %(levelname)7s: %(name)s: %(filename)s:%(lineno)s %(message)s")

    # Color the errors and warnings in red.
    logging.addLevelName(
        logging.ERROR,
        "\033[91m{0}\033[0m".format(logging.getLevelName(logging.ERROR)))
    logging.addLevelName(
        logging.WARNING,
        "\033[91m{0}\033[0m".format(logging.getLevelName(logging.WARNING)))


logger = None


def extend_sytemd_start_timeout_on_cl_vx() -> None:
    """Reference: RM #3571306
    This method specifically deals with Cumulus Linux on VX scenarios
    where the start of the `nvued` SystemD service may take a while
    which extends beyond the default `TimeoutStartSec` service variable.

    See the #3571306 for motivation on why this is being done (useful
    when running QA SSIM2 tests).

    This method checks if the underlying NOS is (Cumulus + VX). This
    won't apply on a hardnode, nor does it run for NVOS or HBN.

    NOTE: TODO: If a similar logic is needed on NVOS, we may have to
                refactor this method to be more generic.
    """
    if not config.NVUE_CL_VX_SYSTEMD_TIMEOUT_EXTEND:
        # If the knob is disabled, do nothing
        return

    if 'WATCHDOG_PID' not in os.environ:
        # Only needed when `nvued` is run as part of SystemD.
        # Not needed when run manually in the foreground.
        return

    import sh

    # First, check if the underlying NOS/platform is Cumulus/VX
    try:
        platform = str(sh.cat("/proc/cmdline", _timeout=1))
        if not("cl_platform" in platform and "cumulus_vx" in platform):
            # If we are not on CL and specifically on a VX,
            # do nothing - For e.g., this logic won't run
            # on a hardnode, OR won't run on NVOS.
            return
    except (sh.ErrorReturnCode, sh.TimeoutException):
        # This may happen either if:
        #  (a) The `/proc/cmdline` file doesn't exist
        #  (b) The `cat /proc/cmdline` command failed (for whatever reason)
        logger.error("Unable to retrieve the NOS/platform type")
        return

    # If we are here, we know that we are running on Cumulus Linux
    # and also running on a VX instance (QCOW2, Vagrant, ...)
    timeout_extension = config.NVUE_CL_VX_SYSTEMD_TIMEOUT_EXTEND_SLICE
    msg = (f"Attempting to extend the SystemD start/stop timeout "
           f"on Cumulus/VX platform to {timeout_extension} minutes")
    logger.info(msg)

    # Convert minutes to microseconds
    timeout_extension_usec = timeout_extension * 60 * 1000000

    try:
        status = [
            f'--status="Extend SystemD timeout to {timeout_extension} mins"',
            f'EXTEND_TIMEOUT_USEC={timeout_extension_usec}'
        ]
        sh.systemd_notify(status, _timeout=1)
    except (sh.ErrorReturnCode, sh.TimeoutException) as err:
        # This may happen either if:
        #  (a) The `systemd-notify` command failed with an error
        #  (b) The `systemd-notify` command execution timed out
        print(str(err))
        logger.error("Failed extending the SystemD start timeout.")
        return

    logger.info("Successfully extended the SystemD start timeout.")


def init():
    global logger

    # Initialize the root logger.
    try:
        loglevel = logging.DEBUG if args.debug else logging.INFO
        if config.NVUE_DEBUG:
            loglevel = logging.DEBUG
        if args.pruning:
            config.NVUE_DYNAMIC_PRUNING = True
        # Set the log level for the root logger.
        logging.getLogger().setLevel(loglevel)

        if args.logfile:
            logger_init_file(loglevel, args.logfile)
        elif args.daemon:
            # In daemon mode, log to syslog by default
            logger_init_syslog(loglevel)
        else:
            # Otherwise, log to stdout by default
            logger_init_stdout(loglevel)

        logger = logging.getLogger()

        logger.info("Starting nvued")

        logger.info(
            "NVUE logging set to %s, chatty deps set to %s",
            logging.getLevelName(loglevel),
            logging.getLevelName(logging.INFO)
        )

        # Some of our deps are a bit too chatty in DEBUG
        logging.getLogger("connexion").setLevel(logging.INFO)
        if loglevel == logging.DEBUG:
            logging.getLogger("sh").setLevel(logging.INFO)
        else:
            logging.getLogger("sh").setLevel(logging.WARN)

        # We don't want profiling to be enabled globally
        # so for now let's restrict this only to nvue modules we
        # are currently profiling. Currently using logger
        # to carry this information around, but could probably use
        # a global variable instead.
        PROFILE_LEVEL_NUM = logging.DEBUG - 1
        logging.PROFILE = PROFILE_LEVEL_NUM
        logging.addLevelName(PROFILE_LEVEL_NUM, "PROFILE")
        if args.profile or config.NVUE_PROFILE:  # pragma: no cover
            logging.getLogger("core.infra.cue.events").setLevel(logging.PROFILE)
            logging.getLogger("core.infra.cue.middleware").setLevel(logging.PROFILE)

    except Exception as e:
        exit("Unable to set up logging:\n{0}\n".format(e))

    # Extend the SystemD start timeout (when/as needed)
    extend_sytemd_start_timeout_on_cl_vx()

    # Make sure we have a working directory
    if os.path.isdir(config.WORKING_DIRECTORY):
        # The working directory already exists, but somebody could
        # have changed its permissions. Run chmod to be sure.
        os.chmod(config.WORKING_DIRECTORY,
                 config.WORKING_DIRECTORY_PERMISSIONS)
    else:
        try:
            os.makedirs(config.WORKING_DIRECTORY,
                        mode=config.WORKING_DIRECTORY_PERMISSIONS)
        except OSError as e:
            logger.error("Cannot create working directory")
            exit(str(e))

    # Tell main it's OK to import units now.
    main.prepare()


if args.logfile == '-' and args.daemon:
    # DaemonContext eats stdout.  That's something we might be able to fix, but
    # for now, don't pretend it's going to work.
    exit("Cannot log to stdout in daemon mode")


# Start it up
if args.daemon:
    # Run cue in the background.
    context = daemon.DaemonContext(
        working_directory=config.WORKING_DIRECTORY,
        signal_map={
            signal.SIGTERM: main.stop,
            signal.SIGINT: main.stop,
            signal.SIGHUP: main.restart,
            signal.SIGUSR2: main.memleak_dump,
        },
        umask=0o22
    )

    context.open()
    with context:
        init()
        main.start()

else:
    init()

    # Run cue in the foreground.
    signal.signal(signal.SIGINT, main.stop)
    signal.signal(signal.SIGTERM, main.stop)
    signal.signal(signal.SIGHUP, main.restart)
    signal.signal(signal.SIGUSR2, main.memleak_dump)
    main.start()


# TODO: FIXME: RM #3033011
# When RM #2923895 and #3017189 were fixed, the fixes only addressed
# automatic enabling and starting of the `nvued` and the `nvue-startup`
# SystemD services when the `python3-nvue` package is (a) Upgraded OR
# (b) Freshly installed. However, if the customer does
# `apt purge python3-nvue` or `apt remove python3-nvue`, the `nvued`
# service is not being stopped gracefully. The SystemD signals
# (SIGTERM/SIGKILL) doesn't reach `nvued` and eventually, the Watchdog
# timer expires, resulting in a SIGABORT being sent to `nvued` which
# terminates nvued abruptly. This is okay since we don't expect customers
# to be uninstalling the NVUE package. Will be revisited later.
# Also see the `faulthandler.enable()` and other `faulthandler` related
# python3 library functions that NVUE uses.
