#!/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 re
import sys

import lib.args
import lib.base
import lib.db_mysql
import lib.human
import lib.txt
from lib.globals import STATE_CRIT, STATE_OK, STATE_UNKNOWN, STATE_WARN

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

DESCRIPTION = """Checks index sizes, fragmentation, and consistent engine and collation usage
across all schemas in MySQL/MariaDB, and lists the largest tables by combined data and index
size so storage growth can be traced before raising memory settings such as the InnoDB buffer
pool. Alerts on mixed storage engines or collations within a single schema, and on table sizes
that cross the optional --warning / --critical thresholds. `--match` and `--ignore` narrow
every aggregate and every check down to a single schema or table. Supports extended
reporting via --lengthy."""

DEFAULT_CRIT = ''
DEFAULT_DEFAULTS_FILE = '/var/spool/icinga2/.my.cnf'
DEFAULT_DEFAULTS_GROUP = 'client'
DEFAULT_LENGTHY = False
DEFAULT_TIMEOUT = 3
DEFAULT_TOP = 10
DEFAULT_WARN = ''

# System schemas: skipped during all aggregates and per-schema checks. mysqltuner
# excludes the same set; `percona` is added because the Percona Server Toolkit installs a
# `percona` schema for its own bookkeeping and it is not user data.
SYSTEM_SCHEMAS = ('information_schema', 'mysql', 'percona', 'performance_schema', 'sys')


def parse_args():
    """Parse command line arguments using argparse."""
    parser = argparse.ArgumentParser(
        description=DESCRIPTION,
        epilog=lib.args.epilog(__file__),
        formatter_class=lib.args.HelpFormatter,
    )

    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 for the size of a single table (data + index). '
        'Supports Nagios ranges with size qualifiers, '
        'for example `10G`, `5G:`, `@1G:10G`. '
        'Default: report only (no alerting).',
        dest='CRITICAL',
        default=DEFAULT_CRIT,
    )

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

    # Append parameters must always use `default=None` instead of a list,
    # because argparse appends user values to the default list instead of
    # replacing it (see https://bugs.python.org/issue16399).
    # After parsing, assign the actual defaults in main() if the value is
    # still None.
    parser.add_argument(
        '--ignore',
        help='Ignore tables whose name matches this Python regular expression. '
        'Matched against the fully qualified table identifier `schema.table`, '
        'so one pattern can drop a whole schema or a single table. Excluded '
        'tables contribute to no aggregate and to no check; system schemas are '
        'skipped unconditionally. '
        'Case-sensitive by default; use `(?i)` for case-insensitive matching. '
        'Can be specified multiple times. '
        'Default: %(default)s. '
        'Example: `--ignore="^(icinga_director|icingaweb2|icingadb)\\."` to skip the '
        'schemas that mix utf8 / utf8mb4 collations by design. '
        'Example: `--ignore="\\.(tmp_|backup_)"` to mute noisy temporary and backup '
        'tables that legitimately differ from the schema-wide engine or collation.',
        dest='IGNORE',
        action='append',
        default=None,
    )

    # Deprecated parameters: hidden from --help, still accepted so existing
    # service definitions keep working. Both are matched against the bare
    # schema resp. table name, which is what they did when MySQL evaluated
    # them via `NOT REGEXP`.
    parser.add_argument(
        '--ignore-schemas',
        help=argparse.SUPPRESS,
        dest='IGNORE_SCHEMAS',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--ignore-tables',
        help=argparse.SUPPRESS,
        dest='IGNORE_TABLES',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--lengthy',
        help=lib.args.help('--lengthy'),
        dest='LENGTHY',
        action='store_true',
        default=DEFAULT_LENGTHY,
    )

    parser.add_argument(
        '--match',
        help='Only check tables whose name matches this Python regular expression. '
        'Matched against the fully qualified table identifier `schema.table`. '
        'A schema without any table cannot match, so it drops out of the report '
        'while this is given. '
        'Case-sensitive by default; use `(?i)` for case-insensitive matching. '
        'Can be specified multiple times. '
        + lib.args.MATCH_IGNORE_PRECEDENCE
        + ' Default: %(default)s. '
        'Example: `--match="^shop\\."` to check the `shop` schema only. '
        'Example: `--match="^shop\\.orders$"` to check one table only.',
        dest='MATCH',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--no-perfdata',
        help=lib.args.help('--no-perfdata'),
        dest='NO_PERFDATA',
        action='store_true',
        default=False,
    )

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

    parser.add_argument(
        '--top',
        help='Number of largest tables (by data + index size) to list. '
        'Default: %(default)s',
        dest='TOP',
        type=int,
        default=DEFAULT_TOP,
    )

    parser.add_argument(
        '-w',
        '--warning',
        help='WARN threshold for the size of a single table (data + index). '
        'Supports Nagios ranges with size qualifiers, '
        'for example `10G`, `5G:`, `@1G:10G`. '
        'Default: report only (no alerting).',
        dest='WARNING',
        default=DEFAULT_WARN,
    )

    args, _ = parser.parse_known_args()
    return args


def get_schemas(conn, excluded):
    # SCHEMA_NAME and others must be uppercase due to MySQL 8+
    sql = f"""
        select SCHEMA_NAME
        from information_schema.schemata
        where SCHEMA_NAME not in ({excluded});
    """  # nosec B608
    return [row['SCHEMA_NAME'] for row in lib.base.coe(lib.db_mysql.select(conn, sql))]


def get_tables(conn, excluded):
    # One row per table for every user schema, aggregated in Python further
    # down. `--match` and `--ignore` are Python regexes and are evaluated
    # here, not by the server, so the aggregates have to be built from the
    # surviving rows. The scan is a once-a-day check.
    sql = f"""
        select TABLE_SCHEMA as table_schema,
            TABLE_NAME as table_name,
            TABLE_TYPE as table_type,
            ENGINE as engine,
            TABLE_COLLATION as table_collation,
            TABLE_ROWS as table_rows,
            DATA_LENGTH as data_length,
            INDEX_LENGTH as index_length
        from information_schema.tables
        where TABLE_SCHEMA not in ({excluded});
    """  # nosec B608
    return lib.base.coe(lib.db_mysql.select(conn, sql))


def get_column_charsets(conn, excluded):
    # `select distinct` collapses this to a handful of rows per table
    # instead of one row per column, which keeps the result set close to
    # the size of the table list itself.
    sql = f"""
        select distinct TABLE_SCHEMA as table_schema,
            TABLE_NAME as table_name,
            CHARACTER_SET_NAME as character_set_name,
            COLLATION_NAME as collation_name
        from information_schema.columns
        where TABLE_SCHEMA not in ({excluded});
    """  # nosec B608
    return lib.base.coe(lib.db_mysql.select(conn, sql))


def compile_filters(args):
    """Compile the filter regexes once. lib.txt.compile_regex() returns a
    (success, result) tuple per pattern and names the parameter in its error
    message, so an invalid pattern exits UNKNOWN via lib.base.coe().
    """
    return {
        key: [lib.base.coe(item) for item in lib.txt.compile_regex(patterns, key)]
        for key, patterns in (
            ('--ignore', args.IGNORE),
            ('--ignore-schemas', args.IGNORE_SCHEMAS),
            ('--ignore-tables', args.IGNORE_TABLES),
            ('--match', args.MATCH),
        )
    }


def keep_table(schema, table, filters):
    """Decide whether a table survives the filters.

    `--match` (include) is applied first, then `--ignore` (exclude), so a
    table hit by `--ignore` is dropped even if it also matches `--match`.
    Both are matched against the fully qualified `schema.table` name, so a
    single pattern addresses either level.

    The deprecated `--ignore-schemas` / `--ignore-tables` are matched
    against the bare schema resp. table name, which is what they did while
    MySQL evaluated them via `NOT REGEXP`.
    """
    if any(item.search(schema) for item in filters['--ignore-schemas']):
        return False
    if any(item.search(table) for item in filters['--ignore-tables']):
        return False
    identifier = f'{schema}.{table}'
    if filters['--match'] and not any(
        item.search(identifier) for item in filters['--match']
    ):
        return False
    return not any(item.search(identifier) for item in filters['--ignore'])


def filter_tables(rows, filters):
    """Drop the rows whose `table_schema` / `table_name` the filters
    exclude.
    """
    return [
        row
        for row in rows
        if keep_table(row['table_schema'] or '', row['table_name'] or '', filters)
    ]


def group_by_schema(rows):
    """Bucket table or column rows by their schema name."""
    grouped = {}
    for row in rows:
        grouped.setdefault(row['table_schema'], []).append(row)
    return grouped


def sum_or_none(values):
    """Sum the non-NULL values. Mirrors SQL `sum()`: a group whose values
    are all NULL sums to NULL, not to 0, and the index-vs-data comparison
    below depends on that distinction.
    """
    present = [int(value) for value in values if value is not None]
    return sum(present) if present else None


def summarize_schema(rows):
    """Aggregate one schema's table rows the way the former per-schema
    `group by TABLE_SCHEMA` query did. Views are counted as tables and
    contribute NULL to the sums, exactly as they did in SQL.
    """
    totals = [
        row['data_length'] + row['index_length']
        if row['data_length'] is not None and row['index_length'] is not None
        else None
        for row in rows
    ]
    return {
        'sum_rows': sum_or_none(row['table_rows'] for row in rows),
        'sum_data': sum_or_none(row['data_length'] for row in rows),
        'sum_index': sum_or_none(row['index_length'] for row in rows),
        'sum_data_index': sum_or_none(totals),
        'cnt_tables': len(rows),
        'engines': sorted({row['engine'] for row in rows if row['engine']}),
        'table_collations': sorted(
            {row['table_collation'] for row in rows if row['table_collation']}
        ),
    }


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

    # logic taken from mysqltuner.pl:mysql_databases(), verified in sync with
    # MySQLTuner (the per-database engine/collation/charset consistency
    # checks and the index-vs-data-size check are unchanged upstream since the
    # original port).
    #
    # Intentional deviation: the index-vs-data-size check additionally requires
    # one of the two sizes to exceed 10 MB. Tiny schemas (under a few MB)
    # routinely have proportionally larger indices than data, which is not
    # actionable for an admin and would generate constant noise.

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

    # set default values for append parameters that were not specified
    if args.IGNORE is None:
        args.IGNORE = []
    if args.IGNORE_SCHEMAS is None:
        args.IGNORE_SCHEMAS = []
    if args.IGNORE_TABLES is None:
        args.IGNORE_TABLES = []
    if args.MATCH is None:
        args.MATCH = []

    # 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, 'SELECT'))

    excluded = ', '.join(f'"{s}"' for s in SYSTEM_SCHEMAS)
    schema_names = get_schemas(conn, excluded)
    all_table_rows = get_tables(conn, excluded)
    all_column_rows = get_column_charsets(conn, excluded)

    lib.db_mysql.close(conn)

    # init some vars
    filters = compile_filters(args)
    # Filter before aggregating, so every schema summary, every consistency
    # finding, the top-tables list and the perfdata describe the same set of
    # tables.
    schema_names = [
        name
        for name in schema_names
        if not any(item.search(name) for item in filters['--ignore-schemas'])
    ]
    tables_by_schema = group_by_schema(filter_tables(all_table_rows, filters))
    columns_by_schema = group_by_schema(filter_tables(all_column_rows, filters))
    schemas_with_tables = {row['table_schema'] for row in all_table_rows}
    # Three cases per schema: it still has tables (scanned), it never had any
    # (empty), or the filters took all of them away (not reported at all). A
    # schema without a table cannot match `--match`, so it drops out of the
    # report while `--match` is given.
    scanned_schemas = []
    empty_schemas = []
    for schema_name in schema_names:
        if schema_name in tables_by_schema:
            scanned_schemas.append(schema_name)
        elif schema_name not in schemas_with_tables and not args.MATCH:
            empty_schemas.append(schema_name)
    state = STATE_OK
    perfdata = ''
    results = {
        'index': '',
        'engine': '',
        'collation': '',
        'colcharset': '',
        'colcollation': '',
    }
    per_schema_rows = []
    total_data_size = 0
    total_index_size = 0
    total_rows = 0
    total_tables = 0
    # Table-size thresholds are Nagios ranges with size qualifiers (e.g. `10G`),
    # converted to byte ranges so lib.base.get_state() can evaluate them. None
    # means "report only" (no alerting).
    warn_bytes = lib.human.humanrange2bytes(args.WARNING) if args.WARNING else None
    crit_bytes = lib.human.humanrange2bytes(args.CRITICAL) if args.CRITICAL else None

    # analyze data
    for schema_name in scanned_schemas:
        summary = summarize_schema(tables_by_schema[schema_name])

        # MySQL returns SQL NULL as Python None (not the string "NULL"); use is/is not
        # None for the guards.
        sum_data = summary['sum_data']
        sum_index = summary['sum_index']
        sum_rows = summary['sum_rows']
        sum_total = summary['sum_data_index']
        cnt_tables = summary['cnt_tables']

        total_data_size += int(sum_data or 0)
        total_index_size += int(sum_index or 0)
        total_rows += int(sum_rows or 0)
        total_tables += int(cnt_tables or 0)

        if (
            sum_data is not None
            and sum_index is not None
            and (sum_data > 10 * 1024 * 1024 or sum_index > 10 * 1024 * 1024)
            and sum_data < sum_index
        ):
            idx_h = lib.human.bytes2human(sum_index)
            data_h = lib.human.bytes2human(sum_data)
            results['index'] += f'{schema_name} ({idx_h} / {data_h}), '

        engines_list = summary['engines']
        tcoll_list = summary['table_collations']
        column_rows = columns_by_schema.get(schema_name, [])
        ccharset_list = sorted(
            {
                row['character_set_name']
                for row in column_rows
                if row['character_set_name']
            }
        )
        ccoll_list = sorted(
            {row['collation_name'] for row in column_rows if row['collation_name']}
        )

        if len(engines_list) > 1:
            results['engine'] += f'{schema_name} ({len(engines_list)}x), '

        if len(tcoll_list) > 1:
            results['collation'] += f'{schema_name} ({len(tcoll_list)}x), '

        if len(ccharset_list) > 1:
            results['colcharset'] += f'{schema_name} ({len(ccharset_list)}x), '

        if len(ccoll_list) > 1:
            results['colcollation'] += f'{schema_name} ({len(ccoll_list)}x), '

        # Compact issue summary for the default (non-lengthy) table.
        issues = []
        if len(engines_list) > 1:
            issues.append(f'{len(engines_list)} engines')
        if len(tcoll_list) > 1:
            issues.append(f'{len(tcoll_list)} table collations')
        if len(ccharset_list) > 1:
            issues.append(f'{len(ccharset_list)} column charsets')
        if len(ccoll_list) > 1:
            issues.append(f'{len(ccoll_list)} column collations')

        # Collect per-schema details for the optional --lengthy table.
        per_schema_rows.append(
            {
                'schema': schema_name,
                'tables': cnt_tables or 0,
                'rows': sum_rows or 0,
                'data': lib.human.bytes2human(sum_data or 0),
                'index': lib.human.bytes2human(sum_index or 0),
                'total': lib.human.bytes2human(sum_total or 0),
                'engines': ', '.join(engines_list) or '-',
                'table_coll': ', '.join(tcoll_list) or '-',
                'col_charsets': ', '.join(ccharset_list) or '-',
                'col_coll': ', '.join(ccoll_list) or '-',
                'issues': ', '.join(issues) if issues else 'OK',
            }
        )

    # Top N tables by combined data + index size, across all scanned schemas.
    # This surfaces cleanup candidates before an admin blindly raises memory
    # settings such as innodb_buffer_pool_size. The list is reported by default;
    # it only drives the state when --warning / --critical size thresholds are
    # given.
    base_tables = [
        row
        for schema_name in scanned_schemas
        for row in tables_by_schema[schema_name]
        if row['table_type'] == 'BASE TABLE'
    ]
    # `coalesce()` semantics: a NULL size counts as 0 here, unlike the
    # per-schema sums above. Sorted by schema and table name as well, so the
    # order of equally sized tables does not jump between runs.
    base_tables.sort(
        key=lambda row: (
            -(int(row['data_length'] or 0) + int(row['index_length'] or 0)),
            row['table_schema'],
            row['table_name'],
        )
    )
    top_tables = base_tables[: max(int(args.TOP), 0)]

    top_rows = []
    for tbl in top_tables:
        data_size = int(tbl['data_length'] or 0)
        index_size = int(tbl['index_length'] or 0)
        total_size = data_size + index_size
        # Thresholds apply to the combined footprint (the table is ranked by it),
        # so the per-row state and its marker live on the Total column.
        tbl_state = lib.base.get_state(
            total_size,
            warn_bytes,
            crit_bytes,
            _operator='range',
        )
        state = lib.base.get_worst(state, tbl_state)
        top_rows.append(
            {
                'schema': tbl['table_schema'],
                'table': tbl['table_name'],
                'data': lib.human.bytes2human(data_size),
                'index': lib.human.bytes2human(index_size),
                'total': f'{lib.human.bytes2human(total_size)}'
                f'{lib.base.state2str(tbl_state, prefix=" ")}',
            }
        )
        # Per-table perfdata so Grafana can trend each table's footprint. Sanitize
        # the `<schema>_<table>` label so exotic names stay perfdata-safe.
        label = (
            re.sub(r'\W+', '_', f'{tbl["table_schema"]}_{tbl["table_name"]}') + '_size'
        )
        perfdata += lib.base.get_perfdata(
            label,
            total_size,
            uom='B',
            warn=warn_bytes,
            crit=crit_bytes,
            _min=0,
        )

    # build the message
    # Independent sections joined with a single blank line between each. This
    # guarantees no `\n\n\n` (double blank lines) appear in the final output
    # regardless of which optional sections fire.

    # Section 1: header + bullet findings, or the OK summary
    bullets = []
    if results['index']:
        bullets.append(
            f'* Index size is larger than data size: {results["index"][:-2]}'
        )
    if results['engine']:
        bullets.append(
            f'* Mixed storage engines (use one engine for all tables in a schema):'
            f' {results["engine"][:-2]}'
        )
    if results['collation']:
        bullets.append(
            f'* Mixed table collations (use one collation for all tables in a schema):'
            f' {results["collation"][:-2]}'
        )
    if results['colcharset']:
        bullets.append(
            f'* Mixed column charsets (use one charset for all text-like columns'
            f' if possible): {results["colcharset"][:-2]}'
        )
    if results['colcollation']:
        bullets.append(
            f'* Mixed column collations (use one collation for all text-like columns'
            f' if possible): {results["colcollation"][:-2]}'
        )

    if bullets:
        state = lib.base.get_worst(state, STATE_WARN)

    # "Everything is ok." leads so admins see the verdict first; the
    # scanned-scope details follow. Format mirrors mysqltuner's per-database
    # stats but condensed to a single line. The state may already be WARN/CRIT
    # from the table-size thresholds above, so the header reflects the worst
    # state, not just the consistency findings.
    summary = (
        f'{len(scanned_schemas) + len(empty_schemas)} user schema(s) scanned,'
        f' {total_tables} table(s),'
        f' {lib.human.number2human(total_rows)} rows,'
        f' {lib.human.bytes2human(total_data_size)} data,'
        f' {lib.human.bytes2human(total_index_size)} indices.'
    )

    sections = []
    if state == STATE_OK:
        sections.append(f'Everything is ok. {summary}')
    else:
        header = (
            'There are critical errors.'
            if state == STATE_CRIT
            else 'There are warnings.'
        )
        if bullets:
            sections.append(f'{header}\n\n' + '\n'.join(bullets))
        else:
            sections.append(header)

    # Section 2: top tables by size. Always reported (the headline feature),
    # independent of --lengthy. The per-row STATE marker is the last thing on
    # the line so IcingaWeb's icon substitution does not break the table.
    if top_rows:
        thresholds = ''
        if args.WARNING or args.CRITICAL:
            thresholds = f' (warn={args.WARNING or "-"} crit={args.CRITICAL or "-"})'
        top_table = lib.base.get_table(
            top_rows,
            ['schema', 'table', 'data', 'index', 'total'],
            header=['Schema', 'Table', 'Data', 'Index', 'Total'],
        )
        sections.append(
            f'Top {len(top_rows)} tables by size{thresholds}:\n\n{top_table}'
        )

    # Section 3: empty-schemas info (common: lazy-init apps, fresh installs,
    # migration leftovers); does not change state, mirroring our "alert on the
    # actionable" stance.
    if empty_schemas:
        sections.append(
            f'Note: {len(empty_schemas)} empty schema(s) (no tables):'
            f' {", ".join(empty_schemas)}.'
        )

    # Section 4: per-schema breakdown table. Without --lengthy: compact summary
    # (Schema | Tables | Total Size | Issues). With --lengthy: full breakdown,
    # mirroring the per-database information mysqltuner emits when invoked with
    # `--dbstat`.
    if per_schema_rows:
        if args.LENGTHY:
            keys = [
                'schema',
                'tables',
                'rows',
                'data',
                'index',
                'total',
                'engines',
                'table_coll',
                'col_charsets',
                'col_coll',
            ]
            headers = [
                'Schema',
                'Tables',
                'Rows',
                'Data',
                'Index',
                'Total',
                'Engines',
                'Table Collations',
                'Column Charsets',
                'Column Collations',
            ]
        else:
            keys = ['schema', 'tables', 'total', 'issues']
            headers = ['Schema', 'Tables', 'Size', 'Issues']
        sections.append(lib.base.get_table(per_schema_rows, keys, header=headers))

    # rstrip each section: lib.base.get_table() appends a trailing newline, which
    # would turn the blank line between two adjacent tables into two blank lines.
    msg = '\n\n'.join(section.rstrip() for section in sections)

    perfdata += lib.base.get_perfdata(
        'mysql_database_count',
        len(scanned_schemas) + len(empty_schemas),
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'mysql_total_data_size',
        total_data_size,
        uom='B',
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'mysql_total_index_size',
        total_index_size,
        uom='B',
        _min=0,
    )
    perfdata += lib.base.get_perfdata('mysql_total_rows', total_rows, _min=0)
    perfdata += lib.base.get_perfdata('mysql_total_tables', total_tables, _min=0)

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


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