#!/opt/vexor/api/venv/bin/python
"""Vexor native check_mssql_health — pymssql reimplementation of the popular
ConSol Labs ``check_mssql_health`` plugin.

Many op5/Nagios command definitions call ``$USER1$/check_mssql_health
--server H --username U --password P --mode MODE ...``. The original is Perl
and needs DBI + DBD::Sybase + FreeTDS, which are painful to ship. This native
version speaks to SQL Server through ``pymssql`` (already in the Vexor API
venv) and implements the commonly-used SQL Server modes with compatible flags,
output and Nagios thresholds. Unimplemented niche modes return UNKNOWN with a
clear message instead of failing to resolve.

Flags (superset of what the imported commands use):
  --server/--hostname H   SQL Server host (required)
  --port P                TCP port (default 1433)
  --username U            login (required)
  --password P            password (required)
  --mode MODE             check mode (see MODES below)
  --name NAME             database / object name for per-object modes
  --warning W             warning threshold (Nagios range)
  --critical C            critical threshold (Nagios range)
  --units U               accepted, informational only
  --regexp                accepted, informational only
  --timeout N             connect/query timeout seconds (default 30)
Accepted and ignored for compatibility: --method --commit --database
  --currentdb --offline --nooffline --report --labelformat --morphmessage ...
"""
from __future__ import annotations

import argparse
import sys
import time

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

try:
    import pymssql
except ImportError as exc:  # pragma: no cover
    print(f"UNKNOWN - pymssql not installed: {exc}")
    sys.exit(UNKNOWN)


# --- Nagios threshold handling ---------------------------------------------
def _range_breach(spec: str, lower_is_bad: bool):
    """Return predicate(value)->bool (True == breach) for a Nagios range.

    For "lower is bad" metrics (hit ratios, free space, page life expectancy)
    a bare number ``N`` is interpreted the check_mssql_health way, i.e. as
    ``N:`` (alert when the value drops below N)."""
    spec = (spec or "").strip()
    if spec == "":
        return lambda v: False
    inside = spec.startswith("@")
    if inside:
        spec = spec[1:]
    if ":" not in spec and not inside and lower_is_bad:
        spec = spec + ":"
    lo, hi = 0.0, float("inf")
    if ":" in spec:
        left, right = spec.split(":", 1)
        lo = float("-inf") if left in ("~", "") else float(left)
        hi = float("inf") if right == "" else float(right)
    else:
        lo, hi = 0.0, float(spec)

    def breach(value: float) -> bool:
        outside = value < lo or value > hi
        return (not outside) if inside else outside

    return breach


def _fmt(v) -> str:
    try:
        f = float(v)
        return str(int(f)) if f.is_integer() else f"{f:.2f}"
    except (TypeError, ValueError):
        return str(v)


def _evaluate(value, warn, crit, label, units, lower_is_bad, text=None):
    """Standard threshold evaluation + perfdata output; returns exit code."""
    perf = f"'{label}'={_fmt(value)}{units};{warn};{crit}"
    msg = text or f"{label} is {_fmt(value)}{units}"
    if crit and _range_breach(crit, lower_is_bad)(float(value)):
        print(f"CRITICAL - {msg} | {perf}")
        return CRITICAL
    if warn and _range_breach(warn, lower_is_bad)(float(value)):
        print(f"WARNING - {msg} | {perf}")
        return WARNING
    print(f"OK - {msg} | {perf}")
    return OK


# --- Performance-counter helpers -------------------------------------------
def _counter(cur, obj_like, counter_like, instance=""):
    cur.execute(
        "SELECT cntr_value FROM sys.dm_os_performance_counters "
        "WHERE RTRIM(object_name) LIKE %s AND RTRIM(counter_name) = %s "
        "AND RTRIM(instance_name) = %s",
        (obj_like, counter_like, instance),
    )
    row = cur.fetchone()
    return None if row is None else float(row[0])


# object suffix, counter name, instance, per-second?, units, lower_is_bad, label
COUNTER_SPECS = {
    "transactions":        ("%:Databases", "Transactions/sec", "_Total", True, "/s", False, "transactions"),
    "batch-requests":      ("%:SQL Statistics", "Batch Requests/sec", "", True, "/s", False, "batch_requests"),
    "sql-initcompilations":("%:SQL Statistics", "SQL Compilations/sec", "", True, "/s", False, "sql_compilations"),
    "sql-recompilations":  ("%:SQL Statistics", "SQL Re-Compilations/sec", "", True, "/s", False, "sql_recompilations"),
    "full-scans":          ("%:Access Methods", "Full Scans/sec", "", True, "/s", False, "full_scans"),
    "lazy-writes":         ("%:Buffer Manager", "Lazy writes/sec", "", True, "/s", False, "lazy_writes"),
    "checkpoint-pages":    ("%:Buffer Manager", "Checkpoint pages/sec", "", True, "/s", False, "checkpoint_pages"),
    "page-life-expectancy":("%:Buffer Manager", "Page life expectancy", "", False, "s", True, "page_life_expectancy"),
    "locks-waits":         ("%:Locks", "Lock Waits/sec", "_Total", True, "/s", False, "lock_waits"),
    "locks-deadlocks":     ("%:Locks", "Number of Deadlocks/sec", "_Total", True, "/s", False, "deadlocks"),
    "locks-timeouts":      ("%:Locks", "Lock Timeouts/sec", "_Total", True, "/s", False, "lock_timeouts"),
    "connected-users":     ("%:General Statistics", "User Connections", "", False, "", False, "connected_users"),
    "total-server-memory": ("%:Memory Manager", "Total Server Memory (KB)", "", False, "", False, "total_server_memory_kb"),
}

RATIO_MODES = {
    "bufferpool-hitrate":               "Buffer cache hit ratio",
    "mem-pool-data-buffer-hit-ratio":   "Buffer cache hit ratio",
}


def _counter_mode(cur, mode, warn, crit):
    obj, cnt, inst, rate, units, lower_bad, label = COUNTER_SPECS[mode]
    v1 = _counter(cur, obj, cnt, inst)
    if v1 is None:
        print(f"UNKNOWN - performance counter '{cnt}' not found")
        return UNKNOWN
    if rate:
        t1 = time.monotonic()
        time.sleep(1.0)
        v2 = _counter(cur, obj, cnt, inst)
        dt = time.monotonic() - t1
        value = max(0.0, (v2 - v1) / dt) if dt > 0 else 0.0
    else:
        value = v1
    return _evaluate(value, warn, crit, label, units, lower_bad)


def _ratio_mode(cur, mode, warn, crit):
    counter = RATIO_MODES[mode]
    cur.execute(
        "SELECT a.cntr_value, b.cntr_value FROM sys.dm_os_performance_counters a "
        "JOIN sys.dm_os_performance_counters b "
        "  ON a.object_name = b.object_name "
        "WHERE RTRIM(a.object_name) LIKE %s AND RTRIM(a.counter_name) = %s "
        "AND RTRIM(b.counter_name) = %s",
        ("%:Buffer Manager", counter, counter + " base"),
    )
    row = cur.fetchone()
    if not row or not row[1]:
        print(f"UNKNOWN - could not read ratio counter '{counter}'")
        return UNKNOWN
    value = float(row[0]) * 100.0 / float(row[1])
    return _evaluate(value, warn, crit, "buffer_cache_hit_ratio", "%", True)


# --- Non-counter modes ------------------------------------------------------
def _mode_connection_time(cur, elapsed, warn, crit):
    return _evaluate(elapsed, warn, crit, "connection_time", "s", False,
                     text=f"connection established in {elapsed:.3f}s")


def _mode_uptime(cur, warn, crit):
    cur.execute("SELECT DATEDIFF(second, sqlserver_start_time, GETDATE()) "
                "FROM sys.dm_os_sys_info")
    secs = float(cur.fetchone()[0])
    days = secs / 86400.0
    return _evaluate(secs, warn, crit, "uptime", "s", True,
                     text=f"server up {days:.1f} days")


def _mode_cpu_busy(cur, warn, crit):
    cur.execute("""
        SELECT TOP 1 100 - record.value(
            '(./Record/SchedulerMonitorEvent/SystemHealth/SystemIdle)[1]', 'int')
        FROM (
            SELECT CONVERT(xml, record) AS record
            FROM sys.dm_os_ring_buffers
            WHERE ring_buffer_type = N'RING_BUFFER_SCHEDULER_MONITOR'
              AND record LIKE '%<SystemHealth>%'
        ) AS x
        ORDER BY record.value('(./Record/@id)[1]', 'bigint') DESC
    """)
    row = cur.fetchone()
    if not row:
        print("UNKNOWN - could not read CPU ring buffer")
        return UNKNOWN
    return _evaluate(float(row[0]), warn, crit, "cpu_busy", "%", False)


def _mode_database_online(cur, name):
    if not name:
        print("UNKNOWN - database-online requires --name")
        return UNKNOWN
    cur.execute("SELECT state_desc FROM sys.databases WHERE name = %s", (name,))
    row = cur.fetchone()
    if not row:
        print(f"CRITICAL - database '{name}' not found")
        return CRITICAL
    state = row[0]
    if state == "ONLINE":
        print(f"OK - database '{name}' is ONLINE")
        return OK
    print(f"CRITICAL - database '{name}' is {state}")
    return CRITICAL


def _db_free_pct(cur, db):
    cur.execute(
        "EXEC('USE [' + %s + ']; SELECT "
        "SUM(CAST(FILEPROPERTY(name, ''SpaceUsed'') AS bigint)) * 8.0 / 1024.0, "
        "SUM(CAST(size AS bigint)) * 8.0 / 1024.0 "
        "FROM sys.database_files WHERE type = 0')",
        (db,),
    )
    row = cur.fetchone()
    if not row or not row[1]:
        return None
    used_mb, alloc_mb = float(row[0]), float(row[1])
    return (alloc_mb - used_mb) / alloc_mb * 100.0


def _mode_database_free(cur, name, warn, crit):
    if name and name.lower() != "all":
        pct = _db_free_pct(cur, name)
        if pct is None:
            print(f"UNKNOWN - could not read free space for '{name}'")
            return UNKNOWN
        return _evaluate(pct, warn, crit, f"free_{name}", "%", True,
                         text=f"database '{name}' has {pct:.1f}% free")
    cur.execute("SELECT name FROM sys.databases WHERE state = 0 AND database_id > 4")
    dbs = [r[0] for r in cur.fetchall()]
    worst = None
    worst_db = None
    for db in dbs:
        try:
            pct = _db_free_pct(cur, db)
        except Exception:
            continue
        if pct is None:
            continue
        if worst is None or pct < worst:
            worst, worst_db = pct, db
    if worst is None:
        print("UNKNOWN - no user databases to check")
        return UNKNOWN
    return _evaluate(worst, warn, crit, "free_min", "%", True,
                     text=f"lowest free space: '{worst_db}' at {worst:.1f}%")


def _mode_backup_age(cur, name, warn, crit, backup_type):
    kind = "log " if backup_type == "L" else ""
    if name and name.lower() != "all":
        cur.execute(
            "SELECT DATEDIFF(hour, MAX(backup_finish_date), GETDATE()) "
            "FROM msdb.dbo.backupset WHERE type = %s AND database_name = %s",
            (backup_type, name),
        )
        row = cur.fetchone()
        age = row[0] if row else None
        if age is None:
            print(f"CRITICAL - database '{name}' has no {kind}backup on record | "
                  f"'{kind.strip() or 'backup'}_age_{name}'=99999h;{warn};{crit}")
            return CRITICAL
        return _evaluate(float(age), warn, crit, f"backup_age_{name}", "h", False,
                         text=f"last {kind}backup of '{name}' was {age}h ago")
    cur.execute(
        "SELECT d.name, DATEDIFF(hour, MAX(b.backup_finish_date), GETDATE()) "
        "FROM sys.databases d "
        "LEFT JOIN msdb.dbo.backupset b "
        "  ON b.database_name = d.name AND b.type = %s "
        "WHERE d.database_id > 4 AND d.state = 0 "
        "GROUP BY d.name",
        (backup_type,),
    )
    worst_age = -1.0
    worst_db = None
    never = []
    for db, age in cur.fetchall():
        if age is None:
            never.append(db)
            worst_age = 99999.0
            worst_db = db
        elif age > worst_age:
            worst_age, worst_db = float(age), db
    if worst_db is None:
        print("OK - no user databases to check")
        return OK
    if never:
        txt = f"never backed up: {', '.join(never[:5])}"
    else:
        txt = f"oldest {kind}backup: '{worst_db}' {int(worst_age)}h ago"
    return _evaluate(worst_age, warn, crit, "backup_age_max", "h", False, text=txt)


def _mode_failed_jobs(cur, warn, crit):
    cur.execute("""
        WITH last AS (
            SELECT job_id, MAX(instance_id) AS mi
            FROM msdb.dbo.sysjobhistory WHERE step_id = 0 GROUP BY job_id)
        SELECT COUNT(*)
        FROM msdb.dbo.sysjobhistory h
        JOIN last ON h.instance_id = last.mi
        JOIN msdb.dbo.sysjobs j ON j.job_id = h.job_id
        WHERE h.run_status = 0 AND j.enabled = 1
    """)
    failed = int(cur.fetchone()[0])
    if not warn and not crit:
        crit = "0"
    return _evaluate(failed, warn, crit, "failed_jobs", "", False,
                     text=f"{failed} SQL Agent job(s) failed on last run")


def _mode_long_running_procs(cur, warn, crit):
    thr = 0
    for spec in (crit, warn):
        try:
            thr = int(float((spec or "").rstrip(":")))
            break
        except ValueError:
            continue
    cur.execute("""
        SELECT COUNT(*)
        FROM sys.dm_exec_requests r
        JOIN sys.dm_exec_sessions s ON r.session_id = s.session_id
        WHERE s.is_user_process = 1 AND r.session_id <> @@SPID
          AND DATEDIFF(second, r.start_time, GETDATE()) > %s
    """, (thr,))
    count = int(cur.fetchone()[0])
    return _evaluate(count, warn or "1", crit or "1", "long_running_procs", "",
                     False, text=f"{count} request(s) running longer than {thr}s")


def _mode_list_databases(cur):
    cur.execute("SELECT name, state_desc FROM sys.databases ORDER BY name")
    rows = cur.fetchall()
    listing = ", ".join(f"{n} ({s})" for n, s in rows)
    print(f"OK - {len(rows)} databases: {listing}")
    return OK


def main() -> int:
    p = argparse.ArgumentParser(prog="check_mssql_health", add_help=True)
    p.add_argument("--server", "--hostname", dest="server", required=True)
    p.add_argument("--port", type=int, default=1433)
    p.add_argument("--username", required=True)
    p.add_argument("--password", required=True)
    p.add_argument("--mode", required=True)
    p.add_argument("--name", default=None)
    p.add_argument("--warning", default="")
    p.add_argument("--critical", default="")
    p.add_argument("--units", default="")
    p.add_argument("--regexp", action="count", default=0)
    p.add_argument("--timeout", type=int, default=30)
    # accepted and ignored for compatibility
    for ign in ("--method", "--database", "--currentdb", "--report",
                "--labelformat", "--morphmessage", "--commit",
                "--offline", "--nooffline", "--extra-opts", "--dbthresholds"):
        p.add_argument(ign, dest=ign.lstrip("-").replace("-", "_"),
                       default=None, nargs="?")
    args, _unknown = p.parse_known_args()

    mode = args.mode.strip()
    warn, crit = args.warning, args.critical

    start = time.monotonic()
    try:
        conn = pymssql.connect(
            server=args.server, user=args.username, password=args.password,
            database="master", port=str(args.port),
            timeout=args.timeout, login_timeout=args.timeout,
        )
    except Exception as exc:  # noqa: BLE001
        print(f"CRITICAL - connection failed: {exc}")
        return CRITICAL
    elapsed = time.monotonic() - start

    try:
        cur = conn.cursor()
        if mode == "connection-time":
            return _mode_connection_time(cur, elapsed, warn or "1", crit or "5")
        if mode == "uptime":
            return _mode_uptime(cur, warn, crit)
        if mode == "cpu-busy":
            return _mode_cpu_busy(cur, warn, crit)
        if mode in COUNTER_SPECS:
            return _counter_mode(cur, mode, warn, crit)
        if mode in RATIO_MODES:
            return _ratio_mode(cur, mode, warn, crit)
        if mode == "database-online":
            return _mode_database_online(cur, args.name)
        if mode == "database-free":
            return _mode_database_free(cur, args.name, warn or "10:", crit or "5:")
        if mode == "database-backup-age":
            return _mode_backup_age(cur, args.name, warn, crit, "D")
        if mode == "database-logbackup-age":
            return _mode_backup_age(cur, args.name, warn, crit, "L")
        if mode == "failed-jobs":
            return _mode_failed_jobs(cur, warn, crit)
        if mode == "long-running-procs":
            return _mode_long_running_procs(cur, warn, crit)
        if mode == "list-databases":
            return _mode_list_databases(cur)
        print(f"UNKNOWN - mode '{mode}' is not implemented by the Vexor native "
              f"check_mssql_health")
        return UNKNOWN
    except Exception as exc:  # noqa: BLE001
        print(f"UNKNOWN - check_mssql_health error in mode '{mode}': {exc}")
        return UNKNOWN
    finally:
        conn.close()


if __name__ == "__main__":
    try:
        sys.exit(main())
    except SystemExit:
        raise
    except Exception as exc:  # noqa: BLE001
        print(f"UNKNOWN - check_mssql_health fatal: {exc}")
        sys.exit(UNKNOWN)
