#!/opt/vexor/api/venv/bin/python
"""Vexor native check_sql — op5/Nagios ``check_sql`` compatible MSSQL check.

Drop-in replacement for the legacy Perl ``check_sql`` (DBD::Sybase/FreeTDS)
that many op5 command definitions reference as ``$USER1$/check_sql``. This
implementation talks to Microsoft SQL Server via ``pymssql`` (already shipped
in the Vexor API venv), so it needs no Perl, FreeTDS or DBD::Sybase.

Supported options (subset used by imported op5 MSSQL commands):
  -H/--hostname   SQL server host (required)
  -U/--username   login (required)
  -P/--password   password (required)
  -D/--database   database (default: master)
  -p/--port       TCP port (default: 1433)
  -q/--query      SQL query to run
  -f/--filename   file containing the SQL query (alternative to -q)
  -e/--expect     expected value (first cell of first row)
  -r/--regexp     treat --expect as a (Python) regular expression
  -W/--rwarning   warning threshold for the RETURNED numeric value
  -C/--rcritical  critical threshold for the RETURNED numeric value
  -w/--warning    warning threshold for the response TIME (seconds)
  -c/--critical   critical threshold for the response TIME (seconds)
  -s/--show       show the query result in the status text
  -l/--label      label for the result (default: Result)
  -d/--driver     accepted and ignored (always MSSQL/pymssql)
  -T/--type       TDS version — accepted and ignored
  -X/--hostconnect accepted and ignored
  -t/--timeout    connection/query timeout in seconds (default: 30)
"""
from __future__ import annotations

import argparse
import re
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)


def _parse_range(spec: str):
    """Return a predicate(value)->bool that is True when value is OUTSIDE the
    acceptable range, following Nagios threshold format.

    Supported: ``N``, ``N:``, ``~:N``, ``M:N``, ``@M:N`` (inside range alerts).
    """
    spec = (spec or "").strip()
    if spec == "":
        return lambda v: False
    inside = spec.startswith("@")
    if inside:
        spec = spec[1:]
    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(value) -> str:
    try:
        f = float(value)
        return str(int(f)) if f.is_integer() else repr(f)
    except (TypeError, ValueError):
        return str(value)


def main() -> int:
    p = argparse.ArgumentParser(prog="check_sql", add_help=True)
    p.add_argument("-H", "--hostname", required=True)
    p.add_argument("-U", "--username", required=True)
    p.add_argument("-P", "--password", required=True)
    p.add_argument("-D", "--database", default="master")
    p.add_argument("-p", "--port", type=int, default=1433)
    p.add_argument("-q", "--query", default=None)
    p.add_argument("-f", "--filename", default=None)
    p.add_argument("-e", "--expect", default=None)
    p.add_argument("-r", "--regexp", action="count", default=0)
    p.add_argument("-W", "--rwarning", default="")
    p.add_argument("-C", "--rcritical", default="")
    p.add_argument("-w", "--warning", default="")
    p.add_argument("-c", "--critical", default="")
    p.add_argument("-s", "--show", action="count", default=0)
    p.add_argument("-l", "--label", default="Result")
    p.add_argument("-t", "--timeout", type=int, default=30)
    # Accepted for op5 compatibility, intentionally ignored:
    p.add_argument("-d", "--driver", default=None)
    p.add_argument("-T", "--type", default=None)
    p.add_argument("-X", "--hostconnect", action="count", default=0)
    args = p.parse_args()

    query = args.query
    if args.filename:
        try:
            with open(args.filename, "r", encoding="utf-8") as fh:
                query = fh.read()
        except OSError as exc:
            print(f"UNKNOWN - cannot read query file: {exc}")
            return UNKNOWN
    if not query or not query.strip():
        print("UNKNOWN - no SQL query given (use -q or -f)")
        return UNKNOWN

    label = args.label.split(",", 1)[0] if args.label else "Result"

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

    try:
        cur = conn.cursor()
        cur.execute(query)
        try:
            row = cur.fetchone()
        except pymssql.OperationalError:
            row = None  # non-result-set statement
    except Exception as exc:  # noqa: BLE001
        conn.close()
        print(f"CRITICAL - query failed: {exc}")
        return CRITICAL
    finally:
        elapsed = time.monotonic() - start

    conn.close()

    cell = row[0] if row else None
    cell_str = "" if cell is None else str(cell)
    shown = f" [{label}: {cell_str}]" if args.show else ""

    # --- Mode 1: expected string / regex ---------------------------------
    if args.expect is not None:
        if row is None:
            print(f"CRITICAL - query returned no rows{shown}")
            return CRITICAL
        if args.regexp:
            matched = re.search(args.expect, cell_str) is not None
        else:
            matched = cell_str == args.expect
        if matched:
            print(f"OK - {label} matched{shown}")
            return OK
        print(f"CRITICAL - {label} did not match "
              f"'{args.expect}' (got '{cell_str}'){shown}")
        return CRITICAL

    # --- Mode 2: numeric returned-value thresholds (-W/-C) ----------------
    if args.rwarning or args.rcritical:
        if row is None:
            print(f"CRITICAL - query returned no rows{shown}")
            return CRITICAL
        try:
            value = float(cell)
        except (TypeError, ValueError):
            print(f"UNKNOWN - returned value not numeric: '{cell_str}'")
            return UNKNOWN
        perf = f"'{label}'={_fmt(value)};{args.rwarning};{args.rcritical}"
        if args.rcritical and _parse_range(args.rcritical)(value):
            print(f"CRITICAL - {label} is {_fmt(value)} | {perf}")
            return CRITICAL
        if args.rwarning and _parse_range(args.rwarning)(value):
            print(f"WARNING - {label} is {_fmt(value)} | {perf}")
            return WARNING
        print(f"OK - {label} is {_fmt(value)} | {perf}")
        return OK

    # --- Mode 3: response-time thresholds (-w/-c) -------------------------
    if args.warning or args.critical:
        perf = f"time={elapsed:.3f}s;{args.warning};{args.critical}"
        if args.critical and _parse_range(args.critical)(elapsed):
            print(f"CRITICAL - query took {elapsed:.3f}s{shown} | {perf}")
            return CRITICAL
        if args.warning and _parse_range(args.warning)(elapsed):
            print(f"WARNING - query took {elapsed:.3f}s{shown} | {perf}")
            return WARNING
        print(f"OK - query took {elapsed:.3f}s{shown} | {perf}")
        return OK

    # --- Default: just run it --------------------------------------------
    print(f"OK - query executed in {elapsed:.3f}s{shown} | "
          f"time={elapsed:.3f}s")
    return OK


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