#!/usr/bin/env python3
"""check_prometheus - scrape a remote Prometheus/OpenMetrics (or JSON) endpoint.

A self-contained Nagios/Naemon check plugin that pulls an EXTERNAL metric INTO
Vexor as a native check, so it reuses Vexor's scheduling, thresholds, perfdata
graphing, SLA and alerting.

It fetches a metrics endpoint over HTTP(S) and either:

  * parses Prometheus/OpenMetrics text exposition and extracts a named metric
    (optionally filtered by labels, optionally aggregated over matching series),
    or
  * parses a JSON document and extracts a numeric value by dotted path
    (``--json-path data.0.value``).

It then applies Nagios threshold ranges and emits a one-line summary plus
perfdata, with the standard exit codes (0 OK, 1 WARNING, 2 CRITICAL,
3 UNKNOWN). Uses only the Python standard library.
"""
from __future__ import annotations

import argparse
import json
import re
import ssl
import sys
import urllib.error
import urllib.request

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

# metric_name{label="v",label2="v2"} 12.34   (labels optional)
_SERIES_RE = re.compile(
    r'^(?P<name>[a-zA-Z_:][a-zA-Z0-9_:]*)'
    r'(?:\{(?P<labels>[^}]*)\})?'
    r'\s+(?P<value>[^\s#]+)'
    r'(?:\s+[0-9.eE+-]+)?\s*$'
)
_LABEL_RE = re.compile(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*=\s*"((?:\\.|[^"\\])*)"')


def die(code: int, msg: str):
    """Print a Nagios-style status line and exit."""
    prefix = {OK: "OK", WARNING: "WARNING", CRITICAL: "CRITICAL"}.get(code, "UNKNOWN")
    print(f"PROMETHEUS {prefix} - {msg}")
    sys.exit(code)


def parse_float(raw: str) -> float:
    """Parse a Prometheus sample value, tolerating NaN/+Inf/-Inf."""
    low = raw.strip().lower()
    if low in ("nan", "+nan", "-nan"):
        return float("nan")
    if low in ("+inf", "inf"):
        return float("inf")
    if low == "-inf":
        return float("-inf")
    return float(raw)


def unescape_label(value: str) -> str:
    return value.replace('\\"', '"').replace("\\\\", "\\").replace("\\n", "\n")


def parse_labels(blob: str) -> dict:
    out = {}
    for m in _LABEL_RE.finditer(blob or ""):
        out[m.group(1)] = unescape_label(m.group(2))
    return out


def fetch(url: str, timeout: float, insecure: bool, headers: list) -> str:
    """Fetch a URL and return the body text, raising on any failure."""
    ctx = None
    if url.lower().startswith("https"):
        ctx = ssl.create_default_context()
        if insecure:
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
    req = urllib.request.Request(url)
    req.add_header("User-Agent", "check_prometheus/1.0 (Vexor)")
    for h in headers or []:
        if ":" not in h:
            raise ValueError(f"invalid header (expected 'Name: value'): {h!r}")
        name, _, val = h.partition(":")
        req.add_header(name.strip(), val.strip())
    with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
        charset = resp.headers.get_content_charset() or "utf-8"
        return resp.read().decode(charset, errors="replace")


def extract_prometheus(body: str, metric: str, want_labels: dict, agg: str):
    """Return (value, matched_count) for the metric/labels using aggregation."""
    values = []
    for line in body.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        m = _SERIES_RE.match(line)
        if not m or m.group("name") != metric:
            continue
        labels = parse_labels(m.group("labels"))
        if want_labels and any(labels.get(k) != v for k, v in want_labels.items()):
            continue
        try:
            values.append(parse_float(m.group("value")))
        except ValueError:
            continue
    if not values:
        return None, 0
    # Drop NaNs from aggregation but keep count of real samples.
    nums = [v for v in values if v == v]  # NaN != NaN
    if not nums:
        return float("nan"), len(values)
    if agg == "first":
        return nums[0], len(values)
    if agg == "max":
        return max(nums), len(values)
    if agg == "min":
        return min(nums), len(values)
    if agg == "avg":
        return sum(nums) / len(nums), len(values)
    return sum(nums), len(values)  # default: sum


def extract_json_path(body: str, path: str) -> float:
    """Extract a numeric value from JSON by dotted path (supports list indices)."""
    data = json.loads(body)
    cur = data
    for part in path.split("."):
        if isinstance(cur, list):
            try:
                cur = cur[int(part)]
            except (ValueError, IndexError):
                raise KeyError(f"bad list index {part!r} in path")
        elif isinstance(cur, dict):
            if part not in cur:
                raise KeyError(f"missing key {part!r} in path")
            cur = cur[part]
        else:
            raise KeyError(f"cannot descend into {part!r}: value is not object/list")
    if isinstance(cur, bool) or not isinstance(cur, (int, float)):
        # allow numeric strings
        return float(cur)
    return float(cur)


class Range:
    """A parsed Nagios threshold range.

    Supported syntax (subset of the Nagios developer guidelines):
      * ``N``         -> alert if value < 0 or value > N      (plain max)
      * ``N:``        -> alert if value < N                   (min)
      * ``:N`` / ``~:N`` -> alert if value > N                (max, neg-inf low)
      * ``A:B``       -> alert if value < A or value > B      (outside range)
      * ``@A:B``      -> alert if A <= value <= B             (inside range)
      * ``~`` denotes negative infinity as the low end.
    """

    def __init__(self, spec: str):
        self.spec = spec
        self.invert = False
        s = spec.strip()
        if s.startswith("@"):
            self.invert = True
            s = s[1:]
        if ":" in s:
            lo, hi = s.split(":", 1)
            self.lo = float("-inf") if lo in ("", "~") else float(lo)
            self.hi = float("inf") if hi == "" else float(hi)
        else:
            self.lo = 0.0
            self.hi = float(s)

    def alerts(self, value: float) -> bool:
        if value != value:  # NaN always alerts
            return True
        inside = self.lo <= value <= self.hi
        return inside if self.invert else not inside


def evaluate(value, warn, crit):
    if crit and crit.alerts(value):
        return CRITICAL
    if warn and warn.alerts(value):
        return WARNING
    return OK


def fmt_value(value: float) -> str:
    if value != value:
        return "NaN"
    if value in (float("inf"), float("-inf")):
        return "Inf" if value > 0 else "-Inf"
    if value == int(value):
        return str(int(value))
    return repr(round(value, 6))


def safe_label(name: str) -> str:
    name = re.sub(r"\s+", "_", name.strip())
    name = re.sub(r"[^a-zA-Z0-9_:.\-]", "_", name)
    return name or "value"


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="check_prometheus",
        description="Scrape a Prometheus/OpenMetrics (or JSON) endpoint and "
                    "return Nagios OK/WARNING/CRITICAL + perfdata.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Threshold ranges (Nagios):\n"
            "  N        alert if value > N (or < 0)\n"
            "  N:       alert if value < N\n"
            "  :N, ~:N  alert if value > N\n"
            "  A:B      alert if value outside [A,B]\n"
            "  @A:B     alert if value inside [A,B]\n"
            "  ~        means negative infinity\n\n"
            "Examples:\n"
            "  check_prometheus -u https://host/metrics -m up -l job=node -w 1: -c 1:\n"
            "  check_prometheus -u https://host/metrics -m http_requests_total "
            "-l code=500 --agg sum -w 10 -c 100\n"
            "  check_prometheus -u https://host/api.json --json-path data.0.value "
            "-w 80 -c 90\n"
        ),
    )
    p.add_argument("-u", "--url", required=True, help="metrics endpoint URL")
    p.add_argument("-m", "--metric", help="metric name to extract (Prometheus mode)")
    p.add_argument("-l", "--label", action="append", default=[],
                   metavar="key=value",
                   help="only match series whose labels include ALL these pairs "
                        "(repeatable)")
    p.add_argument("--json-path", metavar="a.b.c",
                   help="JSON mode: extract numeric value by dotted path "
                        "(supports list indices, e.g. data.0.value)")
    p.add_argument("--agg", choices=["sum", "max", "min", "avg", "first"],
                   default="sum",
                   help="aggregate multiple matching series (default: sum)")
    p.add_argument("-w", "--warning", help="warning threshold range")
    p.add_argument("-c", "--critical", help="critical threshold range")
    p.add_argument("-t", "--timeout", type=float, default=10.0,
                   help="request timeout in seconds (default: 10)")
    p.add_argument("-k", "--insecure", action="store_true",
                   help="skip TLS certificate verification (self-signed)")
    p.add_argument("-H", "--header", action="append", default=[],
                   metavar="'Name: value'",
                   help="extra request header, e.g. 'Authorization: Bearer ...' "
                        "(repeatable)")
    p.add_argument("--label-name", help="override the perfdata metric label")
    return p


def main(argv=None):
    args = build_parser().parse_args(argv)

    json_mode = bool(args.json_path)
    if not json_mode and not args.metric:
        die(UNKNOWN, "either -m/--metric (Prometheus) or --json-path (JSON) is required")

    want_labels = {}
    for pair in args.label:
        if "=" not in pair:
            die(UNKNOWN, f"invalid -l/--label {pair!r} (expected key=value)")
        k, _, v = pair.partition("=")
        want_labels[k.strip()] = v

    try:
        warn = Range(args.warning) if args.warning else None
        crit = Range(args.critical) if args.critical else None
    except ValueError as e:
        die(UNKNOWN, f"bad threshold range: {e}")

    try:
        body = fetch(args.url, args.timeout, args.insecure, args.header)
    except urllib.error.HTTPError as e:
        die(UNKNOWN, f"HTTP {e.code} fetching {args.url}: {e.reason}")
    except urllib.error.URLError as e:
        die(UNKNOWN, f"cannot reach {args.url}: {e.reason}")
    except Exception as e:  # noqa: BLE001 - any fetch error -> UNKNOWN
        die(UNKNOWN, f"request failed: {e}")

    count = 1
    if json_mode:
        try:
            value = extract_json_path(body, args.json_path)
        except json.JSONDecodeError as e:
            die(UNKNOWN, f"response is not valid JSON: {e}")
        except (KeyError, ValueError, TypeError) as e:
            die(UNKNOWN, f"json-path '{args.json_path}' failed: {e}")
        label = safe_label(args.label_name or args.json_path)
    else:
        value, count = extract_prometheus(body, args.metric, want_labels, args.agg)
        if value is None:
            lbl = (" with labels " + ",".join(f"{k}={v}" for k, v in want_labels.items())
                   ) if want_labels else ""
            die(UNKNOWN, f"metric '{args.metric}'{lbl} not found at {args.url}")
        label = safe_label(args.label_name or args.metric)

    code = evaluate(value, warn, crit)

    vstr = fmt_value(value)
    warn_field = args.warning if args.warning else ""
    crit_field = args.critical if args.critical else ""
    perf = f"{label}={vstr};{warn_field};{crit_field}"
    extra = ""
    if not json_mode and count > 1:
        extra = f" ({args.agg} of {count} series)"
    die(code, f"{label}={vstr}{extra} | {perf}")


if __name__ == "__main__":
    main()
