#! /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 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")
args = parser.parse_args()


def logger_init_syslog(loglevel):
    syslog_h = logging.handlers.SysLogHandler(address="/dev/log")
    formatter = logging.Formatter("nvued: %(levelname)7s:  %(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: %(message)s")
    else:
        logging.basicConfig(
            filename=logfile,
            filemode='a',
            level=loglevel,
            format="%(asctime)s %(levelname)7s: %(name)s: %(message)s")


def logger_init_stdout(loglevel, logfile=None):
    logging.basicConfig(
        level=loglevel,
        format="%(asctime)s %(levelname)7s: %(name)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 init():
    global logger

    # Initialize the root logger.
    try:
        loglevel = logging.DEBUG if args.debug else logging.INFO
        if config.NVUE_DEBUG:
            loglevel = logging.DEBUG

        # 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)

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

    # 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.
