#!/usr/bin/env python3
"""Nagios/Icinga checks for the authenticated RustProxy administration API."""
from __future__ import annotations

import argparse
import http.cookiejar
import json
import os
from pathlib import Path
import ssl
import sys
import time
import urllib.error
import urllib.request

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


class CheckError(RuntimeError):
    pass


def plugin_exit(state: int, message: str, perf: list[str] | None = None) -> None:
    names = {OK: "OK", WARNING: "WARNING", CRITICAL: "CRITICAL", UNKNOWN: "UNKNOWN"}
    suffix = f" | {' '.join(perf)}" if perf else ""
    print(f"RUSTPROXY {names[state]} - {message}{suffix}")
    raise SystemExit(state)


def threshold(value: float, warning: float, critical: float) -> int:
    if value >= critical:
        return CRITICAL
    if value >= warning:
        return WARNING
    return OK


class Client:
    def __init__(self, args: argparse.Namespace) -> None:
        self.base = args.url.rstrip("/")
        self.username = args.user or os.environ.get("RUSTPROXY_NAGIOS_USER")
        self.password = os.environ.get(args.password_env)
        self.cookie_file = Path(args.cookie_file).expanduser() if args.cookie_file else None
        self.jar = http.cookiejar.MozillaCookieJar(str(self.cookie_file)) if self.cookie_file else http.cookiejar.CookieJar()
        if self.cookie_file and self.cookie_file.exists():
            try:
                self.jar.load(ignore_discard=True, ignore_expires=False)
            except (OSError, http.cookiejar.LoadError):
                pass
        if args.insecure:
            context = ssl._create_unverified_context()
        else:
            context = ssl.create_default_context(cafile=args.ca_file)
        self.opener = urllib.request.build_opener(
            urllib.request.HTTPCookieProcessor(self.jar),
            urllib.request.HTTPSHandler(context=context),
        )
        self.timeout = args.timeout

    def _request(self, path: str, payload: dict | None = None) -> tuple[dict, float]:
        body = None if payload is None else json.dumps(payload, separators=(",", ":")).encode()
        request = urllib.request.Request(
            self.base + path,
            data=body,
            method="GET" if payload is None else "POST",
            headers={"Accept": "application/json", "Content-Type": "application/json", "X-RustProxy-Request": "1"},
        )
        started = time.monotonic()
        try:
            with self.opener.open(request, timeout=self.timeout) as response:
                raw = response.read(8 * 1024 * 1024)
        except urllib.error.HTTPError:
            raise
        except (OSError, urllib.error.URLError) as error:
            raise CheckError(str(error)) from error
        elapsed_ms = (time.monotonic() - started) * 1000.0
        try:
            value = json.loads(raw)
        except json.JSONDecodeError as error:
            raise CheckError("administration API returned invalid JSON") from error
        if not isinstance(value, dict):
            raise CheckError("administration API returned an unexpected payload")
        return value, elapsed_ms

    def login(self) -> None:
        if not self.username or self.password is None:
            raise CheckError(
                "no valid session and no credentials; set --user/RUSTPROXY_NAGIOS_USER and the password environment variable"
            )
        try:
            self._request("/api/auth/login", {"username": self.username, "password": self.password})
        except urllib.error.HTTPError as error:
            raise CheckError(f"login failed with HTTP {error.code}") from error
        if self.cookie_file:
            self.cookie_file.parent.mkdir(parents=True, exist_ok=True)
            self.jar.save(ignore_discard=True, ignore_expires=False)
            try:
                os.chmod(self.cookie_file, 0o600)
            except OSError:
                pass

    def status(self) -> tuple[dict, float]:
        try:
            return self._request("/api/status")
        except urllib.error.HTTPError as error:
            if error.code != 401:
                raise CheckError(f"status endpoint returned HTTP {error.code}") from error
        self.login()
        try:
            return self._request("/api/status")
        except urllib.error.HTTPError as error:
            raise CheckError(f"status endpoint returned HTTP {error.code} after login") from error


def pools(status: dict, selected: list[str]) -> list[dict]:
    values = status.get("backend_pools") or []
    if not isinstance(values, list):
        raise CheckError("status payload has no backend_pools array")
    if not selected:
        return values
    wanted = set(selected)
    found = [pool for pool in values if str(pool.get("id", "")) in wanted]
    missing = wanted - {str(pool.get("id", "")) for pool in found}
    if missing:
        raise CheckError("unknown vhost/backend pool: " + ", ".join(sorted(missing)))
    return found


def check_api(client: Client, args: argparse.Namespace) -> None:
    status, latency = client.status()
    state = threshold(latency, args.warning_ms, args.critical_ms)
    revision = status.get("revision", "?")
    uptime = status.get("uptime_seconds", "?")
    plugin_exit(state, f"API {latency:.1f} ms, revision {revision}, uptime {uptime}s", [f"api_ms={latency:.1f}ms;{args.warning_ms};{args.critical_ms};0"])


def check_backends(client: Client, args: argparse.Namespace) -> None:
    status, api_ms = client.status()
    down: list[str] = []
    checking: list[str] = []
    up = ignored = 0
    for pool in pools(status, args.pool):
        pool_id = str(pool.get("id", "?"))
        for member in pool.get("members") or []:
            state = str(member.get("status", "unknown"))
            name = f"{pool_id}/{member.get('id', '?')}"
            if state == "down":
                down.append(name)
            elif state == "checking":
                checking.append(name)
            elif state == "up":
                up += 1
            elif state in {"disabled", "draining"}:
                ignored += 1
            else:
                checking.append(f"{name}({state})")
    state = CRITICAL if down else WARNING if checking else OK
    detail = f"up={up}, down={len(down)}, checking={len(checking)}, ignored={ignored}"
    if down:
        detail += "; down: " + ", ".join(down[:10])
    elif checking:
        detail += "; checking: " + ", ".join(checking[:10])
    plugin_exit(state, detail, [f"up={up}", f"down={len(down)}", f"checking={len(checking)}", f"api_ms={api_ms:.1f}ms"])


def check_latency(client: Client, args: argparse.Namespace) -> None:
    status, api_ms = client.status()
    samples: list[tuple[float, str]] = []
    for pool in pools(status, args.pool):
        pool_id = str(pool.get("id", "?"))
        for member in pool.get("members") or []:
            if str(member.get("status", "")) in {"disabled", "draining"}:
                continue
            raw = member.get("ewma_latency_ms") or member.get("latency_ms") or 0
            try:
                value = float(raw)
            except (TypeError, ValueError):
                continue
            if value > 0:
                samples.append((value, f"{pool_id}/{member.get('id', '?')}"))
    if not samples:
        plugin_exit(WARNING, "no backend latency samples are available", [f"api_ms={api_ms:.1f}ms"])
    value, name = max(samples)
    state = threshold(value, args.warning_ms, args.critical_ms)
    plugin_exit(state, f"slowest backend {name}: {value:.1f} ms ({len(samples)} sampled)", [f"backend_ms={value:.1f}ms;{args.warning_ms};{args.critical_ms};0", f"api_ms={api_ms:.1f}ms"])


def check_queues(client: Client, args: argparse.Namespace) -> None:
    status, api_ms = client.status()
    logs = int(status.get("logs_dropped") or 0)
    cold = int((status.get("cache") or {}).get("cold_queue_dropped") or 0)
    total = logs + cold
    state = CRITICAL if total >= args.critical else WARNING if total >= args.warning else OK
    plugin_exit(state, f"dropped work: logs={logs}, cold-cache={cold}", [f"logs_dropped={logs};{args.warning};{args.critical};0", f"cold_queue_dropped={cold};{args.warning};{args.critical};0", f"api_ms={api_ms:.1f}ms"])


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("mode", choices=["api", "backends", "latency", "queues"])
    p.add_argument("--url", required=True, help="Administration origin, e.g. https://rustproxy-admin.example.net")
    p.add_argument("--user", help="Administrator username; defaults to RUSTPROXY_NAGIOS_USER")
    p.add_argument("--password-env", default="RUSTPROXY_NAGIOS_PASSWORD", help="environment variable containing the password")
    p.add_argument("--cookie-file", help="persistent session cookie jar; recommended for frequent checks")
    p.add_argument("--ca-file", help="custom CA bundle for the administration TLS certificate")
    p.add_argument("--insecure", action="store_true", help="disable TLS verification (diagnostics only)")
    p.add_argument("--timeout", type=float, default=5.0)
    p.add_argument("--pool", action="append", default=[], help="limit backend checks to a vhost/backend-pool id; repeatable")
    p.add_argument("--warning-ms", type=float, default=250.0)
    p.add_argument("--critical-ms", type=float, default=1000.0)
    p.add_argument("--warning", type=int, default=1, help="queues: warning total")
    p.add_argument("--critical", type=int, default=10, help="queues: critical total")
    return p


def main() -> None:
    args = parser().parse_args()
    if args.warning_ms < 0 or args.critical_ms <= args.warning_ms:
        plugin_exit(UNKNOWN, "invalid latency thresholds")
    if args.warning < 0 or args.critical <= args.warning:
        plugin_exit(UNKNOWN, "invalid queue thresholds")
    try:
        client = Client(args)
        {"api": check_api, "backends": check_backends, "latency": check_latency, "queues": check_queues}[args.mode](client, args)
    except CheckError as error:
        plugin_exit(UNKNOWN, str(error))
    except Exception as error:  # Nagios plugins must fail closed with a single line.
        plugin_exit(UNKNOWN, f"unexpected check failure: {error}")


if __name__ == "__main__":
    main()