#!/usr/bin/env python3
"""Checkmk special agent for Proxmox Backup Server (REST, API-token auth)."""
from __future__ import annotations
import argparse
import json
import os
import re
import signal
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import oposs_pbs_collect as collect_mod
from oposs_pbs_cache import StateCache
from oposs_pbs_client import PbsClient, resolve_password_ref

try:  # cmk.utils.password_store is not a stable API path; guard for offline use
    from cmk.utils.password_store import replace_passwords
except ImportError:
    def replace_passwords():
        return None


def parse_args(argv):
    p = argparse.ArgumentParser(description="Proxmox Backup Server special agent")
    p.add_argument("hostaddress")
    p.add_argument("--port", type=int, default=8007)
    p.add_argument("--token-id", required=True)
    p.add_argument("--token-secret", required=True)
    p.add_argument("--no-verify-tls", action="store_true")
    p.add_argument("--cacert")
    p.add_argument("--include-datastore", action="append", default=[])
    p.add_argument("--exclude-datastore", action="append", default=[])
    p.add_argument("--task-limit", type=int, default=1000)
    p.add_argument("--timeout", type=int, default=60)  # per-request HTTP timeout
    # Wall-clock ceiling on expensive /snapshots refreshes per run; 0 = unlimited.
    # Kept below Checkmk's cmc_check_timeout (60s by default): a run that
    # overruns it is killed and prints nothing at all, so every guest of this
    # PBS loses its piggyback record for that run.
    p.add_argument("--refresh-budget", type=int, default=45)
    p.add_argument("--piggyback-template", default="{guest}")
    p.add_argument("--piggyback-regex")  # "PATTERN=REPLACEMENT"
    p.add_argument("--no-piggyback-datastore", action="append", default=[])
    # Repeatable regex; matched (re.search) against "<store>/<ns>/<type>/<id>".
    p.add_argument("--ignore-backup", action="append", default=[])
    # Checkmk host name of this PBS host. A backup group resolving to exactly
    # this name is this server's own backup (PBS running as a VM on the cluster
    # it backs up); it is reported inline instead of as self-referential
    # piggyback.
    p.add_argument("--self-host")
    p.add_argument("--cache-dir")
    p.add_argument("--test-file")
    p.add_argument("--now", type=int)  # test hook; real runs use time.time()
    return p.parse_args(argv)


class _FileClient:
    def __init__(self, routes):
        self._r = routes
    def get(self, path, params=None, timeout=None):
        return self._r.get(path)


def _make_client(args):
    if args.test_file:
        with open(args.test_file, encoding="utf-8") as fh:
            return _FileClient(json.load(fh))
    return PbsClient(args.hostaddress, args.port, args.token_id, args.token_secret,
                     verify=not args.no_verify_tls, cafile=args.cacert,
                     timeout=args.timeout)


def _cache_path(args):
    """Where the learned per-group state lives.

    Under var/, not tmp/: $OMD_ROOT/tmp is a tmpfs that `omd restart` empties.
    This cache is not scratch -- it holds the guest names and cadence history
    the agent has learned, and losing it makes every backup group fall back to
    its bare VMID until the refresh budget has worked through the backlog.
    """
    base = args.cache_dir or os.path.join(
        os.environ.get("OMD_ROOT", "/var/tmp"), "var", "check_mk", "oposs_pbs")
    return os.path.join(base, f"{args.hostaddress}.json")


def _print_section(name, payload):
    print(f"<<<{name}:sep(0)>>>")
    print(json.dumps(payload, separators=(",", ":")))


def main(argv=None):
    replace_passwords()  # best-effort: resolves the legacy --pwstore argv form
    args = parse_args(argv if argv is not None else sys.argv[1:])
    # server_side_calls passes the token as a bare Secret, i.e. an inline
    # "<pw_id>:<pw_store_file>" reference that replace_passwords() leaves
    # untouched; resolve it explicitly to the real secret.
    args.token_secret = resolve_password_ref(args.token_secret)
    regex = None
    if args.piggyback_regex and "=" in args.piggyback_regex:
        pat, repl = args.piggyback_regex.split("=", 1)
        regex = (pat, repl)

    # Compile once, at startup: an invalid pattern is a misconfiguration and
    # must fail loudly rather than silently suppressing nothing (or everything).
    try:
        ignore = [re.compile(pat) for pat in args.ignore_backup]
    except re.error as exc:
        print(f"invalid --ignore-backup pattern {exc.pattern!r}: {exc}",
              file=sys.stderr)
        return 2

    opts = collect_mod.Options(
        include=args.include_datastore, exclude=args.exclude_datastore,
        task_limit=args.task_limit, piggyback_template=args.piggyback_template,
        piggyback_regex=regex, no_piggyback=set(args.no_piggyback_datastore),
        ignore=ignore)

    import time
    now = args.now if args.now is not None else int(time.time())
    cache_path = _cache_path(args)
    cache = StateCache.load(cache_path)
    client = _make_client(args)

    def _save():
        try:
            cache.save(cache_path)
        except OSError:
            pass

    # Persist forward progress if the fetcher kills us mid-collect, so cold
    # datastores warm up over successive runs instead of restarting from zero.
    def _on_term(_signum, _frame):
        _save()
        os._exit(0)
    for _sig in (signal.SIGTERM, signal.SIGINT):
        try:
            signal.signal(_sig, _on_term)
        except (ValueError, OSError):
            pass  # not in main thread (e.g. under test) -> skip

    budget = collect_mod.RefreshBudget(
        args.refresh_budget if args.refresh_budget > 0 else None)
    host, piggyback = collect_mod.collect(client, opts, cache, now,
                                          budget=budget, save=_save)

    _print_section("oposs_pbs_server", host.get("oposs_pbs_server", {}))
    if "oposs_pbs_datastore" in host:
        _print_section("oposs_pbs_datastore", host["oposs_pbs_datastore"])
    if "oposs_pbs_jobs" in host:
        _print_section("oposs_pbs_jobs", host["oposs_pbs_jobs"])
    if "oposs_pbs_backup_rollup" in host:
        _print_section("oposs_pbs_backup_rollup", host["oposs_pbs_backup_rollup"])

    # Split the groups into this host's own backup and everyone else's. The
    # inline section must be printed before the first piggyback marker,
    # otherwise it would be attributed to the preceding piggyback host.
    self_host = (args.self_host or "").casefold()
    own: list = []
    remote: list = []
    for host_name, record in piggyback:
        if not host_name:
            continue
        if self_host and host_name.casefold() == self_host:
            own.append(record)
        else:
            remote.append((host_name, record))

    if own:
        print("<<<oposs_pbs_backup:sep(0)>>>")
        for record in own:
            print(json.dumps(record, separators=(",", ":")))

    for host_name, record in remote:
        print(f"<<<<{host_name}>>>>")
        _print_section("oposs_pbs_backup", record)
        print("<<<<>>>>")

    _save()
    # Always exit 0: an unreachable PBS is surfaced via reachable:False in the
    # oposs_pbs_server section (check reports CRIT), not via the agent exit code.
    return 0


if __name__ == "__main__":
    sys.exit(main())
