#!/usr/bin/env python3
# -*- coding: utf-8; py-indent-offset: 4 -*-
#
# Author:  Linuxfabrik GmbH, Zurich, Switzerland
# Contact: info (at) linuxfabrik (dot) ch
#          https://www.linuxfabrik.ch/
# License: The Unlicense, see LICENSE file.

# https://github.com/Linuxfabrik/monitoring-plugins/blob/main/CONTRIBUTING.md

"""See the check's README for more details."""

import argparse
import sys

import lib.args
import lib.base
import lib.db_mysql
from lib.globals import STATE_OK, STATE_UNKNOWN

__author__ = 'Linuxfabrik GmbH, Zurich/Switzerland'
__version__ = '2026051101'

DESCRIPTION = """Runs up to two admin-supplied SQL SELECT statements against MySQL/MariaDB
and checks each result against a Nagios range expression. One statement is paired
with `--warning` and triggers WARN; the second with `--critical` and triggers CRIT.
A query returning one row with one column is checked as a single value (useful for
`SELECT COUNT(*) ...`, `SELECT MAX(timestamp) ...`, etc.); any other shape is checked
by its row count. Useful for custom application-level monitoring: queue depth, stale
rows, replication lag tables, total user count, daily order count, failed-job count,
inventory below threshold, expiring licences, anything an application stores in a
table."""

DEFAULT_DEFAULTS_FILE = '/var/spool/icinga2/.my.cnf'
DEFAULT_DEFAULTS_GROUP = 'client'
DEFAULT_TIMEOUT = 3

# Maximum rows shown in the result table appended to the plugin output. Beyond
# this we keep only the first and last few rows; the full result is still
# evaluated for state, just not echoed back.
MAX_RESULT_ROWS = 10
RESULT_HEAD = 5
RESULT_TAIL = 5

# Maximum displayed length of the SQL statement in the output line. Long
# queries get an ellipsis so the summary stays one readable line.
MAX_QUERY_DISPLAY = 80


def parse_args():
    """Parse command line arguments using argparse."""
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    parser.add_argument(
        '-V',
        '--version',
        action='version',
        version=f'%(prog)s: v{__version__} by {__author__}',
    )

    parser.add_argument(
        '--always-ok',
        help=lib.args.help('--always-ok'),
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c',
        '--critical',
        help='CRIT threshold as a Nagios range expression.',
        dest='CRIT',
    )

    parser.add_argument(
        '--critical-query',
        help='`SELECT` statement whose result is checked against `--critical`. '
        'If the result contains more than one column, the row count is used. '
        'Otherwise the single returned value is used.',
        dest='CRITICAL_QUERY',
    )

    parser.add_argument(
        '--defaults-file',
        help='MySQL/MariaDB cnf file to read user, host and password from. '
        'Example: `--defaults-file=/var/spool/icinga2/.my.cnf`. '
        'Default: %(default)s',
        dest='DEFAULTS_FILE',
        default=DEFAULT_DEFAULTS_FILE,
    )

    parser.add_argument(
        '--defaults-group',
        help=lib.args.help('--defaults-group') + ' Default: %(default)s',
        dest='DEFAULTS_GROUP',
        default=DEFAULT_DEFAULTS_GROUP,
    )

    parser.add_argument(
        '--timeout',
        help=lib.args.help('--timeout') + ' Default: %(default)s (seconds)',
        dest='TIMEOUT',
        type=int,
        default=DEFAULT_TIMEOUT,
    )

    parser.add_argument(
        '-w',
        '--warning',
        help='WARN threshold as a Nagios range expression.',
        dest='WARN',
    )

    parser.add_argument(
        '--warning-query',
        help='`SELECT` statement whose result is checked against `--warning`. '
        'If the result contains more than one column, the row count is used. '
        'Otherwise the single returned value is used.',
        dest='WARNING_QUERY',
    )

    args, _ = parser.parse_known_args()
    return args


def trunc_query(query):
    """Single-line, ellipsised version of `query` for the summary line."""
    one_line = ' '.join(query.split())
    if len(one_line) > MAX_QUERY_DISPLAY:
        return one_line[: MAX_QUERY_DISPLAY - 3] + '...'
    return one_line


def evaluate_query(conn, query, threshold, kind):
    """Run `query`, derive its value (single cell or row count) and grade it
    against `threshold` for the given `kind` ('warn' or 'crit'). Returns a
    tuple of (state, value, rows, shortened).
    """
    state = STATE_OK
    value = 0
    rows = []
    shortened = False
    if not query:
        return state, value, rows, shortened

    rows = lib.base.coe(lib.db_mysql.select(conn, query))
    if not rows:
        return state, value, rows, shortened

    if len(rows) == 1 and len(rows[0]) == 1:
        # Single-cell result: e.g. `SELECT COUNT(*) FROM ...` or
        # `SELECT MAX(updated_at) FROM ...`. Use the cell's value directly.
        value = next(iter(rows[0].values()))
    else:
        # Multiple rows, or multi-column row: grade by row count, but still
        # echo head + tail of the rows so admins can investigate.
        value = len(rows)
        if len(rows) > MAX_RESULT_ROWS:
            rows = rows[:RESULT_HEAD] + rows[-RESULT_TAIL:]
            shortened = True

    if kind == 'warn':
        state = lib.base.get_state(value, threshold, None, _operator='range')
    else:
        state = lib.base.get_state(value, None, threshold, _operator='range')
    return state, value, rows, shortened


def render_result_table(rows, shortened):
    """Build the optional `head ... tail` table appended below the summary."""
    if not rows:
        return ''
    parts = []
    if shortened:
        parts.append(
            f'\nAttention: Table below is truncated, showing the'
            f' {RESULT_HEAD} first and the {RESULT_TAIL} last rows.'
        )
    try:
        keys = rows[0].keys()
        parts.append('\n' + lib.base.get_table(rows, keys, header=keys))
    except (AttributeError, IndexError):
        return ''
    return ''.join(parts)


def main():
    """The main function. This is where the magic happens."""

    # parse the command line
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    if args.WARNING_QUERY is None and args.CRITICAL_QUERY is None:
        lib.base.cu('Nothing to check, no queries provided.')

    # fetch data
    mysql_connection = {
        'defaults_file': args.DEFAULTS_FILE,
        'defaults_group': args.DEFAULTS_GROUP,
        'timeout': args.TIMEOUT,
    }
    conn = lib.base.coe(lib.db_mysql.connect(mysql_connection))
    lib.base.coe(lib.db_mysql.check_privileges(conn))

    # init some vars
    state = STATE_OK
    msg = ''
    perfdata = ''
    facts = []

    # analyze data
    state_warn, value_warn, rows_warn, shortened_warn = evaluate_query(
        conn, args.WARNING_QUERY, args.WARN, 'warn',
    )
    state = lib.base.get_worst(state, state_warn)
    state_crit, value_crit, rows_crit, shortened_crit = evaluate_query(
        conn, args.CRITICAL_QUERY, args.CRIT, 'crit',
    )
    state = lib.base.get_worst(state, state_crit)

    lib.db_mysql.close(conn)

    # build the message
    if args.WARNING_QUERY:
        facts.append(
            f'WARN query (`{trunc_query(args.WARNING_QUERY)}`)'
            f' returned {value_warn}'
            f'{lib.base.state2str(state_warn, prefix=" ")}'
        )
        perfdata += lib.base.get_perfdata(
            'mysql_query_warn_value',
            value_warn,
            warn=args.WARN,
        )
    if args.CRITICAL_QUERY:
        facts.append(
            f'CRIT query (`{trunc_query(args.CRITICAL_QUERY)}`)'
            f' returned {value_crit}'
            f'{lib.base.state2str(state_crit, prefix=" ")}'
        )
        perfdata += lib.base.get_perfdata(
            'mysql_query_crit_value',
            value_crit,
            crit=args.CRIT,
        )
    if state == STATE_OK:
        # "Everything is ok." leads the OK message so the admin sees the
        # verdict first; the facts follow.
        msg = 'Everything is ok. ' + '. '.join(facts) + '.'
    else:
        msg = '. '.join(facts) + '.'

    # build table output
    msg += render_result_table(rows_warn, shortened_warn)
    msg += render_result_table(rows_crit, shortened_crit)

    # over and out
    lib.base.oao(msg, state, perfdata, always_ok=args.ALWAYS_OK)


if __name__ == '__main__':
    try:
        main()
    except Exception:
        lib.base.cu()
