"""Base class and auxiliary classes for NAGIOS probes."""

from glob import glob
import argparse
from functools import cached_property
from configparser import ConfigParser, InterpolationError, NoOptionError
import logging
import os
from io import StringIO
import sys
import tempfile
import time
import traceback
from typing import Any, Callable, Dict, List, NoReturn, Optional, Tuple
from arcnagios import confargparse
from arcnagios.utils import counted_noun

_module_log = logging.getLogger(__name__)

OK = 0
WARNING = 1
CRITICAL = 2
UNKNOWN = 3

_STATUS_NAME_BY_CODE = ['OK', 'WARNING', 'CRITICAL', 'UNKNOWN']
_STATUS_CODE_BY_NAME = {
    'OK': OK,
    'WARNING': WARNING,
    'CRITICAL': CRITICAL,
    'UNKNOWN': UNKNOWN,
}

def status_name(status):
    return _STATUS_NAME_BY_CODE[status]

def status_by_name(name):
    try:
        return _STATUS_CODE_BY_NAME[name.upper()]
    except KeyError as exn:
        raise ValueError('%s is not a Nagios status name.' % name) from exn

class ServiceReport(Exception):
    def __init__(self, status: int, message: str):
        self.status = status
        self.message = message
        Exception.__init__(self, status_name(status) + ': ' + message)

class ServiceOk(ServiceReport):
    def __init__(self, message: str):
        ServiceReport.__init__(self, OK, message)

class ServiceWarning(ServiceReport):
    def __init__(self, message: str):
        ServiceReport.__init__(self, WARNING, message)

class ServiceCritical(ServiceReport):
    def __init__(self, message: str):
        ServiceReport.__init__(self, CRITICAL, message)

class ServiceUnknown(ServiceReport):
    def __init__(self, message: str):
        ServiceReport.__init__(self, UNKNOWN, message)

class NagiosPerfdata:
    def __init__(
            self, label: str, value: float, uom: Optional[str] = None,
            limit_warn: Optional[float] = None,
            limit_crit: Optional[float] = None,
            limit_min: Optional[float] = None,
            limit_max: Optional[float] = None):
        self.label = label
        self.value = value
        self.uom = uom
        self.limit_warn = limit_warn
        self.limit_crit = limit_crit
        self.limit_min = limit_min
        self.limit_max = limit_max

    def __str__(self):
        limits = [
            self.limit_warn, self.limit_crit,
            self.limit_min,  self.limit_max
        ]
        def to_str(x):
            if x is None:
                return ''
            else:
                return '%g' % x
        return '\'%s\'=%s%s;' % (self.label, self.value, self.uom or '') \
             + ';'.join(map(to_str, limits))

    def __iadd__(self, other):
        assert self.label == other.label
        assert self.uom == other.uom
        self.value += other.value
        if other.limit_warn:
            self.limit_warn += other.limit_warn
        if other.limit_crit:
            self.limit_crit += other.limit_crit
        if other.limit_min:
            self.limit_min += other.limit_min
        if other.limit_max:
            self.limit_max += other.limit_max
        return self

class NagiosPerflog:
    def __init__(self, granular: bool = False):
        self._granular = granular
        self._perfdatas: Dict[str, NagiosPerfdata] = {}

    def add(self, label: str, value: float, uom: Optional[str] = None,
            limit_warn: Optional[float] = None,
            limit_crit: Optional[float] = None,
            limit_min: Optional[float] = None,
            limit_max: Optional[float] = None) -> None:
        perfdata = NagiosPerfdata(
                label, value, uom=uom,
                limit_warn=limit_warn, limit_crit=limit_crit,
                limit_min=limit_min, limit_max=limit_max)
        if label in self._perfdatas:
            perfdata += self._perfdatas[label]
        self._perfdatas[label] = perfdata

    def addi(self, label: str, index: Optional[str],
             value: float, uom: Optional[str] = None,
             limit_warn: Optional[float] = None,
             limit_crit: Optional[float] = None,
             limit_min: Optional[float] = None,
             limit_max: Optional[float] = None) -> None:
        if self._granular and not index is None:
            label += '[' + index + ']'
        self.add(label, value, uom=uom,
                 limit_warn=limit_warn, limit_crit=limit_crit,
                 limit_min=limit_min, limit_max=limit_max)

    def is_empty(self) -> bool:
        return self._perfdatas == {}

    def __str__(self) -> str:
        return ' '.join(map(str, self._perfdatas.values()))

class NagiosPerflogTimer:

    def __init__(self, perflog: NagiosPerflog,
                 label: str, index: Optional[str]):
        self._perflog: NagiosPerflog = perflog
        self._label: str = label
        self._index: Optional[str] = index
        self._start_time: Optional[float] = None
        self._stop_time: Optional[float] = None

    def __enter__(self):
        self._start_time = time.time()

    def __exit__(self, exc_type, exc_value, exc_tb):
        if self._start_time:
            self._stop_time = time.time()
            elapsed = self._stop_time - self._start_time
            self._perflog.addi(self._label, self._index,
                               elapsed, uom='s', limit_min=0)

class NagiosReport:
    """Instances of this class collects information to be reported to Nagios.
    You should use one instance for the active check result, and one instance
    for each passive (``host_name``, ``service_description``) kombination to
    target."""

    def __init__(self, host_name: str, service_description: str):
        self._log_buffer = StringIO()
        self.host_name = host_name
        self.service_description = service_description
        self.log = logging.Logger('%s/%s'%(host_name, service_description))
        self.log.addHandler(logging.StreamHandler(self._log_buffer))
        self.status_code = OK
        self.status_messages: List[List[str]] = [[], [], [], []]
        self.status_code_counts: List[int] = [0, 0, 0, 0]

    def update_status_code(self, status_code: int):
        """Update the Nagios exit code to the maximum of ``status_code`` and
        the current code."""

        assert OK <= status_code <= UNKNOWN
        self.status_code_counts[status_code] += 1
        if status_code > self.status_code:
            self.status_code = status_code

    def update_status(
            self, status_code: int, status_message: Optional[str] = None) \
            -> None:
        """Update the Nagios exit code to the maximum of ``status_code`` and
        the current code, and add ``status_message`` to the messages to be
        used in case no higher exit codes overrides this call."""

        self.update_status_code(status_code)
        if status_message:
            self.status_messages[status_code].append(status_message)

    def status_message(self, subject: Optional[str] = None) -> str:
        """Format a status message suitable as the first line of output to
        Nagios.  This will be based on the calls to `update_status` which are
        relevant for the final exit code.  If no messages are registered for
        the code, make a generic message, possibly referring to `subject`."""

        if self.status_messages[self.status_code] == []:
            name = subject or 'service'
            count = self.status_code_counts[self.status_code]
            if self.status_code == OK:
                return '%s OK'%(name.capitalize())
            if self.status_code == WARNING:
                return '%s in %s'%(counted_noun(count, 'warning'), name)
            if self.status_code == CRITICAL:
                return '%s in %s'%(counted_noun(count, 'error'), name)
            return 'Failed to check %s'%name
        return ' '.join(self.status_messages[self.status_code])

    @property
    def status_details(self) -> str:
        """A string containing messages logged to `self.log`."""

        return self._log_buffer.getvalue()

_LOG_LEVEL_BY_NAME = {
    'DEBUG': logging.DEBUG,
    'INFO': logging.INFO,
    'WARN': logging.WARN,
    'WARNING': logging.WARNING,
    'ERROR': logging.ERROR,
    'CRITICAL': logging.CRITICAL,
}

class NagiosPlugin:

    probe_name: str
    bugtracker_url: str = 'http://bugzilla.nordugrid.org/'
    main_config_section: List[str] = []

    def __init__(
            self, use_host: bool = False, use_port: bool = False,
            default_port: Optional[int] = None):
        # Set up a logger for collecting messages and initialize perfdata.
        self.perflog: NagiosPerflog
        self._passive_reports: Dict[Tuple[str, str], NagiosReport] = {}

        def scan_loglevel(s):
            try:
                return int(s)
            except ValueError:
                pass
            try:
                return _LOG_LEVEL_BY_NAME[s.upper()]
            except KeyError:
                # pylint: disable=W0707
                raise argparse.ArgumentTypeError('Invalid loglevel.')

        # Define the base argument parsers. Plug-in implementations will
        # extend it or the sub-parsers.
        self.argparser = confargparse.ConfigArgumentParser()
        argp = self.argparser.add_argument_group('General Options')
        argp.add_argument('--loglevel', dest = 'loglevel', type = scan_loglevel,
                default = 'WARNING',
                help = 'Set the log level for NAGIOS messages. '
                       'Use either numeric or symbolic names corresponding to '
                       'the definitions in the python logging module.')
        seems_like_manual = os.getenv('NAGIOS_HOSTNAME') is None
        argp.add_argument('--how-invoked', dest = 'how_invoked',
                choices = ['manual', 'nagios'],
                default = seems_like_manual and 'manual' or 'nagios',
                help = 'Indicate how invoked. '
                       '(Default: Check a Nagios enviroment variable).')
        argp.add_argument('--command-file', dest = 'command_file',
                default = os.getenv('NAGIOS_COMMANDFILE',
                                    '/var/spool/nagios/cmd/nagios.cmd'),
                help = 'The Nagios command file.  By default the '
                       '$NAGIOS_COMMANDFILE environment variable is used, '
                       'which is usually what you want.')
        argp.add_argument('--multiline-separator',
                dest = 'multiline_separator', default = ' - ',
                help = 'Replacement for newlines when submitting multiline '
                       'results to passive services.  Pass the empty string '
                       'drop extra lines.')
        argp.add_argument('--clip-passive-status', type=int,
                default = None,
                help = 'The maximum number of characters to publish to a '
                       'passive service status, including the ellipses which '
                       'will be added to indicate the that it is clipped.')
        if use_host:
            argp.add_argument('-H', '--host', dest = 'host',
                    help = 'Host of service to check.')
        if use_port:
            argp.add_argument('-p', '-P', '--port', dest = 'port', type = int,
                    default = default_port,
                    help = 'Port number of service to connect to.  '
                           'OBS! The capital -P alternative is deprecated.')
        argp.add_argument('--dump-options', dest = 'dump_options',
                default = False, action = 'store_true',
                help = 'Dump options to standard output.  '
                       'Believed to be useful for debugging.')
        argp.add_argument('--arcnagios-spooldir', dest = 'arcnagios_spooldir',
                default = '/var/spool/arc/nagios',
                help = 'Top-level spool directory to be used by the ARC '
                       'Nagios plugins.')
        argp.add_argument('--home-dir', dest = 'home_dir',
                help = 'Override $HOME at startup.  This is a workaround '
                       'for external commands which store things under '
                       '$HOME on systems where the nagios account does '
                       'not have an appropriate or writable home directory.')
        argp.add_argument('--import', dest = 'imports', action = 'append',
                default = [], metavar = 'MODULE',
                help = 'Import the given module.  The module initializer can '
                       'register implementations for arcnagios.substitution '
                       'or arcnagios.jobplugins, or modify the enviroment.')
        self._remove_at_exit: List[str] = []
        self._at_exit: List[Callable[[], None]] = []
        self._created_tmpdir: bool = False
        self.opts: Any = None
        self.nagios_report: NagiosReport
        self.perflog = NagiosPerflog()
        self.log: logging.Logger

        # Initialize mix-ins if any.
        super().__init__()

    def config_dir(self) -> str:
        return os.getenv('ARCNAGIOS_CONFIG_DIR', '/etc/arc/nagios')

    def config_paths(self) -> List[str]:
        """A list of paths to search for the configuration."""

        config_paths_env = os.getenv('ARCNAGIOS_CONFIG')
        if not config_paths_env is None:
            return config_paths_env.split(':')
        else:
            config_paths = glob(os.path.join(self.config_dir(), '*.ini'))
            config_paths.sort()
            return config_paths

    @cached_property
    def config(self) -> ConfigParser:
        environment = {}
        environment['epoch_time'] = str(int(time.time()))
        parser = ConfigParser(defaults = environment)
        parser.read(self.config_paths())
        return parser

    def at_exit(self, f: Callable[[], None]) -> None:
        self._at_exit.append(f)

    def remove_at_exit(self, *paths: str) -> None:
        self._remove_at_exit += paths

    def tmpdir(self) -> str:
        tmpdir = os.path.join(self.opts.arcnagios_spooldir, 'tmp')
        if not self._created_tmpdir and not os.path.exists(tmpdir):
            try:
                os.makedirs(tmpdir)
            except OSError as exn:
                self.nagios_exit(UNKNOWN, 'Cannot create %s: %s'%(tmpdir, exn))
        return tmpdir

    def mkstemp(self, suffix: str = '', prefix: str = 'tmp') -> Tuple[int, str]:
        tmpdir = self.tmpdir()
        fd, path = tempfile.mkstemp(suffix, prefix, tmpdir)
        self.remove_at_exit(path)
        return fd, path

    def mktemp(self, suffix: str = '', prefix: str = 'tmp') -> str:
        tmpdir = self.tmpdir()
        path = tempfile.mktemp(suffix, prefix, tmpdir)
        self.remove_at_exit(path)
        return path

    def parse_args(self, args: List[str]) -> None:
        try:
            if self.main_config_section:
                self.argparser.configure_defaults(
                        self.config, self.main_config_section)
            self.opts = self.argparser.parse_args(args)
        except Exception as exn: # pylint: disable=broad-except
            # TODO: Reconsider, we should rather consider it invalid.
            self.nagios_report = NagiosReport('UNKNOWN', 'ACTIVE')
            self.log = self.nagios_report.log
            raise exn

        host_name = getattr(self.opts, 'host', '')
        self.nagios_report = NagiosReport(host_name, 'ACTIVE')
        self.log = self.nagios_report.log
        if hasattr(self.opts, 'granular_perfdata'):
            self.perflog = NagiosPerflog(self.opts.granular_perfdata)

        if not self.opts.home_dir is None:
            os.environ['HOME'] = self.opts.home_dir

        # Set the log level.  The level of self.log determines how much is
        # reported as additional lines to Nagios.
        self.log.setLevel(self.opts.loglevel)
        if self.opts.how_invoked == 'manual':
            # File-level loggers are used for debugging.  Only increase the
            # level if manually invoked, since we may otherwise obscure the
            # input to Nagios.
            _module_log.root.setLevel(self.opts.loglevel)

        for module_name in self.opts.imports:
            try:
                __import__(module_name)
            except ImportError as exn:
                raise ServiceUnknown('Error importing %s: %s'
                                     % (module_name, exn)) from exn

    def check(self) -> Optional[ServiceReport]:
        """Override this method in your plugin to implement the test.  You
        should not call this directly, but use `nagios_run`, which will handle
        argument parsing and report results to Nagios."""

        raise NotImplementedError('The `check` method has not been implemented.')

    def nagios_run(self) -> NoReturn:
        """This is the method to invoke from the probe script.  It parses
        command-line options, calls the probe, prints Nagios-formatted output
        to stdout, and exits with an appropriate code."""

        # Parse command-line arguments and configuration file.
        try:
            self.parse_args(sys.argv[1:])
        except confargparse.UsageError as exn:
            self.log.error(str(exn))
            self.log.error('Need --help?')
            self.nagios_exit(UNKNOWN, 'Invalid command-line options: %s' % exn)
        except confargparse.ConfigError as exn:
            self.log.error(str(exn))
            self.nagios_exit(UNKNOWN, 'Invalid configuration: %s' % exn)
        if self.opts.dump_options:
            for k, v in vars(self.opts).items():
                self.log.info('%s = %r', k, v)

        # Run the metric and report.
        try:
            service_report = self.check()
        except ServiceReport as exn:
            service_report = exn
        except SystemExit as exn:
            raise exn
        except NoOptionError as exn:
            self.nagios_exit(UNKNOWN, str(exn))
        except InterpolationError as exn:
            for line in str(exn).split('\n'):
                self.log.error(line)
            self.nagios_exit(UNKNOWN, 'Error in configuration file.')
        except Exception as exn: # pylint: disable=broad-except
            _, _, trace = sys.exc_info()
            self.log.error('----%<----')
            self.log.error('Please report this bug to %s including '
                           'the following:', self.bugtracker_url)
            self.log.error('%r', exn)
            self.log.error('Traceback:')
            for line in traceback.format_tb(trace):
                self.log.error(str(line.strip()))
            self.log.error('----%<----')
            self.nagios_exit(UNKNOWN, 'Bug in Nagios probe.')
        if service_report is None:
            self.nagios_exit()
        if not isinstance(service_report, ServiceReport):
            msg = 'Invalid value %r returned by plugin.' % service_report
            self.nagios_exit(UNKNOWN, msg)
        self.nagios_exit(service_report.status, service_report.message)

    def add_perfdata(
            self, label: str, value: float, uom: Optional[str] = None,
            limit_warn: Optional[float] = None,
            limit_crit: Optional[float] = None,
            limit_min: Optional[float] = None,
            limit_max: Optional[float] = None
        ) -> None:
        self.perflog.add(
                label, value, uom=uom,
                limit_warn=limit_warn,
                limit_crit=limit_crit,
                limit_min=limit_min,
                limit_max=limit_max)

    def nagios_report_for(
            self, host_name: str, service_description: str, create = True):
        """Return a `NagiosReport` instance which will be submitted to the
        passive service `service_description` on `host_name`.  Each time you
        call this with the same `host_name` and `service_description`, you
        will get the same instance."""

        key = (host_name, service_description)
        report = self._passive_reports.get(key, None)
        if report is None and create:
            report = NagiosReport(host_name, service_description)
            self._passive_reports[key] = report
        return report

    def nagios_exit(
            self, status_code: int = OK, status_message: Optional[str] = None,
            subject: Optional[str] = None) -> NoReturn:
        """Submit all passives check results, update `self.nagios_report` with
        the given `status_code` and `status_message` if specified, then
        communicate `self.nagios_report` as the active check result.  In
        particular, this writes out status messages, perfdatas and logging
        written to `self.nagios_report`, and calls `sys.exit` with
        `self.nagios_report.status_code`."""

        for f in self._at_exit:
            f()

        for path in self._remove_at_exit:
            try:
                os.remove(path)
            except OSError:
                pass

        for report in self._passive_reports.values():
            self.submit_passive_service_result(
                    report.host_name,
                    report.service_description,
                    report.status_code,
                    report.status_message(subject),
                    report.status_details)

        self.nagios_report.update_status(status_code, status_message)
        sys.stdout.write(self.nagios_report.status_message(subject))
        if self.perflog and not self.perflog.is_empty():
            sys.stdout.write('|' + str(self.perflog))
        sys.stdout.write('\n')
        sys.stdout.write(self.nagios_report.status_details)
        sys.exit(self.nagios_report.status_code)

    def submit_passive_service_result(
            self, host_name: str, svc_description: str,
            return_code: int, plugin_output: str, details = None):
        """Manually submit a passive service result.  It is in general more
        convenient to use `nagios_report_for` and let `nagios_run` submit the
        results."""

        t_chk = int(time.time())
        if details and self.opts.multiline_separator:
            # FIXME:  How do we submit multi-line results to passive services?
            # Currently it seems not to be supported.  This will look ugly,
            # but it's better than leaving the operator clueless.
            sep = self.opts.multiline_separator
            plugin_output += sep + details.strip().replace('\n', sep)
        if self.opts.clip_passive_status \
                and len(plugin_output) > self.opts.clip_passive_status:
            plugin_output = plugin_output[:self.opts.clip_passive_status-4] \
                    .rsplit(None, 1)[0] + ' ...'
        rstr = '[%d] PROCESS_SERVICE_CHECK_RESULT;%s;%s;%d;%s\n' % \
                (t_chk, host_name, svc_description, return_code, plugin_output)
        if self.opts.how_invoked == 'manual':
            _module_log.info(rstr)
            return
        try:
            with open(self.opts.command_file, 'w', encoding='utf-8') as fd:
                fd.write(rstr)
        except IOError as exn:
            raise ServiceUnknown('Cannot write to NAGIOS command file %s: %s'
                                 % (self.opts.command_file, exn)) from exn
