#!/usr/bin/env python3
# Copyright (C) 2026 Christian Wirtz <doc@snowheaven.de>
# License: GPL-2.0-only, see LICENSE in the repository root.
"""
mkp-cleanup - find and clean up outdated / duplicate Checkmk extension packages (MKPs)

Scans the output of the site-local `mkp` command line tool for the two areas
also shown in the Checkmk GUI under "Setup > Maintenance > Extension packages":

  * "Enabled (inactive on this site)"  - packages that are enabled but not
    actually active on this site (e.g. because they don't match the running
    Checkmk version).
  * "All packages (enabled or disabled)" - every installed package,
    additionally checked for duplicate/superseded versions of the same
    package, and optionally (--include-until-version-exceeded) for packages
    past their stated "Until Version".

It never touches the currently active, newest version of a package - only
packages that are already inactive, superseded by a newer installed
version, or (opt-in) past their stated "Until Version" are ever considered
candidates.

Must be run as the Checkmk site user (i.e. `mkp` must be reachable on PATH,
$OMD_ROOT/$OMD_SITE must be set - e.g. via `omd su <site>` or
`su - <site>`).

See `mkp-cleanup --help` and the README for details and examples.
"""

from __future__ import annotations

import argparse
import datetime
import json
import os
import re
import subprocess
import sys
from dataclasses import dataclass, field

PROG = "mkp-cleanup"
VERSION = "1.0.0"

STATE_ACTIVE = "Enabled (active on this site)"
STATE_INACTIVE = "Enabled (inactive on this site)"
STATE_DISABLED = "Disabled"

REASON_INACTIVE = "inactive-on-site"
REASON_SUPERSEDED = "superseded-duplicate"
REASON_UNTIL_EXCEEDED = "until-version-exceeded"

HEADER_LABELS = [
    "Name",
    "Version",
    "Title",
    "Author",
    "Req. Version",
    "Until Version",
    "Files",
    "State",
]


class MkpError(RuntimeError):
    """Raised whenever the `mkp` command line tool fails or its output can't be parsed."""


# --------------------------------------------------------------------------- #
# data model
# --------------------------------------------------------------------------- #

@dataclass
class Package:
    name: str
    version: str
    title: str
    author: str
    req_version: str
    until_version: str
    files: str
    state: str
    reasons: list = field(default_factory=list)

    @property
    def is_candidate(self) -> bool:
        return bool(self.reasons)

    @property
    def key(self):
        return (self.name, self.version)


# --------------------------------------------------------------------------- #
# environment / safety
# --------------------------------------------------------------------------- #

def check_site_context(skip_check: bool) -> None:
    if skip_check:
        return
    omd_site = os.environ.get("OMD_SITE")
    omd_root = os.environ.get("OMD_ROOT")
    if not omd_site or not omd_root:
        raise MkpError(
            "This does not look like a Checkmk site shell (OMD_SITE/OMD_ROOT "
            "not set). Run this as the site user, e.g.:\n"
            "  omd su <sitename>\n"
            "or\n"
            "  su - <sitename>\n"
            "Use --skip-site-check to bypass this guard (e.g. for testing)."
        )


def find_mkp_binary() -> str:
    from shutil import which

    path = which("mkp")
    if not path:
        raise MkpError(
            "`mkp` was not found on PATH. Make sure you are running this "
            "inside a Checkmk site shell."
        )
    return path


def get_current_cmk_version() -> str | None:
    """Best-effort detection of the running Checkmk version, e.g. '2.3.0p10'."""
    omd_root = os.environ.get("OMD_ROOT")
    if not omd_root:
        return None
    version_link = os.path.join(omd_root, "version")
    try:
        target = os.path.realpath(version_link)
    except OSError:
        return None
    base = os.path.basename(target)
    # strip edition suffix, e.g. "2.3.0p10.cee" -> "2.3.0p10"
    base = re.sub(r"\.(cre|cee|cme|cfe|cse|cpe)$", "", base)
    return base or None


# --------------------------------------------------------------------------- #
# mkp invocation
# --------------------------------------------------------------------------- #

def run_mkp(args: list, mkp_bin: str) -> str:
    try:
        result = subprocess.run(
            [mkp_bin] + args,
            check=False,
            capture_output=True,
            text=True,
        )
    except OSError as exc:
        raise MkpError(f"failed to execute `mkp {' '.join(args)}`: {exc}") from exc
    if result.returncode != 0:
        raise MkpError(
            f"`mkp {' '.join(args)}` exited with {result.returncode}: "
            f"{result.stderr.strip() or result.stdout.strip()}"
        )
    return result.stdout


def fetch_packages(mkp_bin: str) -> list:
    output = run_mkp(["list"], mkp_bin)
    return parse_mkp_list(output)


def parse_mkp_list(output: str) -> list:
    lines = [l for l in output.splitlines() if l.strip()]
    if not lines:
        return []

    header_line = None
    header_idx = 0
    for idx, line in enumerate(lines):
        if "Name" in line and "State" in line:
            header_line = line
            header_idx = idx
            break
    if header_line is None:
        raise MkpError(
            "could not recognize the output of `mkp list` (no 'Name ... "
            "State' header found). Your Checkmk version's mkp output format "
            "may not be supported - please check `mkp list` manually."
        )

    positions = []
    search_from = 0
    try:
        for label in HEADER_LABELS:
            pos = header_line.index(label, search_from)
            positions.append(pos)
            search_from = pos + len(label)
    except ValueError as exc:
        raise MkpError(
            f"unexpected `mkp list` header layout: {header_line!r} ({exc})"
        ) from exc

    packages = []
    for line in lines[header_idx + 1:]:
        if re.fullmatch(r"[-=\s]+", line):
            # separator line under the header (e.g. "--------  -------  ...")
            continue
        padded = line.ljust(positions[-1] + 1)
        fields = []
        for i, start in enumerate(positions):
            end = positions[i + 1] if i + 1 < len(positions) else len(padded)
            fields.append(padded[start:end].strip())
        name, version, title, author, req_version, until_version, files, state = fields
        if not name:
            continue
        packages.append(
            Package(
                name=name,
                version=version,
                title=title,
                author=author,
                req_version=req_version,
                until_version=until_version,
                files=files,
                state=state,
            )
        )
    return packages


def fetch_show(mkp_bin: str, name: str, version: str) -> str:
    try:
        return run_mkp(["show", name, version], mkp_bin)
    except MkpError as exc:
        return f"  (could not fetch details: {exc})"


# --------------------------------------------------------------------------- #
# version comparison
# --------------------------------------------------------------------------- #

def version_key(version: str):
    """Best-effort, dependency-free version sort key.

    Splits the version string into alternating digit / non-digit groups so
    that e.g. '2' < '10' (unlike plain string comparison) while still
    tolerating arbitrary suffixes like '2.3.0p10' or '1.0-beta1'.
    """
    parts = re.findall(r"\d+|\D+", version or "")
    key = []
    for part in parts:
        if part.isdigit():
            key.append((1, int(part), ""))
        else:
            key.append((0, 0, part))
    return key


# --------------------------------------------------------------------------- #
# classification
# --------------------------------------------------------------------------- #

def classify(
    packages: list,
    keep_latest: int,
    current_version: str | None,
    include_until_version_exceeded: bool = False,
) -> None:
    by_name: dict = {}
    for pkg in packages:
        by_name.setdefault(pkg.name, []).append(pkg)

    for pkg in packages:
        if pkg.state == STATE_INACTIVE:
            pkg.reasons.append(REASON_INACTIVE)

        if (
            include_until_version_exceeded
            and current_version
            and pkg.until_version
            and pkg.until_version.lower() not in ("none", "")
            and re.search(r"\d", pkg.until_version)
            and version_key(current_version) > version_key(pkg.until_version)
        ):
            pkg.reasons.append(f"{REASON_UNTIL_EXCEEDED} ({pkg.until_version})")

    for name, versions in by_name.items():
        if len(versions) <= keep_latest:
            continue
        ordered = sorted(versions, key=lambda p: version_key(p.version), reverse=True)
        for pkg in ordered[keep_latest:]:
            newest = ordered[0].version
            pkg.reasons.append(f"{REASON_SUPERSEDED} (newer version {newest} installed)")


# --------------------------------------------------------------------------- #
# selection
# --------------------------------------------------------------------------- #

def select_enabled_inactive(packages: list, pattern) -> list:
    return [p for p in packages if p.state == STATE_INACTIVE and matches(p, pattern)]


def select_all_candidates(packages: list, pattern) -> list:
    return [p for p in packages if p.is_candidate and matches(p, pattern)]


def matches(pkg: Package, pattern) -> bool:
    if pattern is None:
        return True
    return bool(pattern.search(pkg.name))


# --------------------------------------------------------------------------- #
# output helpers
# --------------------------------------------------------------------------- #

class Colors:
    def __init__(self, enabled: bool):
        self.enabled = enabled

    def _wrap(self, code: str, text: str) -> str:
        return f"\033[{code}m{text}\033[0m" if self.enabled else text

    def bold(self, text): return self._wrap("1", text)
    def red(self, text): return self._wrap("31", text)
    def yellow(self, text): return self._wrap("33", text)
    def green(self, text): return self._wrap("32", text)
    def dim(self, text): return self._wrap("2", text)


def print_table(headers: list, rows: list) -> None:
    if not rows:
        print("  (none)")
        return
    widths = [len(h) for h in headers]
    for row in rows:
        for i, cell in enumerate(row):
            widths[i] = max(widths[i], len(str(cell)))
    fmt = "  ".join("{:<%d}" % w for w in widths)
    print("  " + fmt.format(*headers))
    print("  " + fmt.format(*["-" * w for w in widths]))
    for row in rows:
        print("  " + fmt.format(*[str(c) for c in row]))


def candidate_rows(candidates: list) -> list:
    return [
        (p.name, p.version, p.state, ", ".join(p.reasons))
        for p in sorted(candidates, key=lambda p: (p.name, version_key(p.version)))
    ]


# --------------------------------------------------------------------------- #
# logging
# --------------------------------------------------------------------------- #

def default_log_path() -> str:
    omd_root = os.environ.get("OMD_ROOT")
    if omd_root:
        return os.path.join(omd_root, "var", "log", "mkp-cleanup.log")
    return os.path.join(os.getcwd(), "mkp-cleanup.log")


def log_line(log_file: str, message: str) -> None:
    ts = datetime.datetime.now().isoformat(timespec="seconds")
    user = os.environ.get("USER", "?")
    line = f"{ts} | {user} | {message}\n"
    try:
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        with open(log_file, "a", encoding="utf-8") as fh:
            fh.write(line)
    except OSError as exc:
        print(f"warning: could not write log file {log_file}: {exc}", file=sys.stderr)


# --------------------------------------------------------------------------- #
# cleanup execution
# --------------------------------------------------------------------------- #

def do_cleanup(
    candidates: list,
    mkp_bin: str,
    dry_run: bool,
    purge: bool,
    log_file: str,
    colors: Colors,
) -> tuple:
    disabled, removed, failed, skipped = 0, 0, 0, 0
    for pkg in candidates:
        if pkg.state == STATE_ACTIVE:
            # Defensive guard: classify() should never mark the sole active
            # version as a candidate, but never act on it if it somehow is -
            # regardless of --purge.
            print(colors.yellow(f"  skipping {pkg.name} {pkg.version}: currently active, refusing to touch"))
            skipped += 1
            continue

        actions = []
        if pkg.state != STATE_DISABLED:
            actions.append(("disable", [pkg.name, pkg.version]))
        if purge:
            actions.append(("remove", [pkg.name, pkg.version]))

        if not actions:
            # Already Disabled and no --purge: there is nothing left to do.
            print(colors.dim(f"  {pkg.name} {pkg.version}: already disabled, nothing to do without --purge"))
            log_line(log_file, f"SKIP (already disabled, no --purge) {pkg.name} {pkg.version}")
            skipped += 1
            continue

        for verb, args in actions:
            label = f"{verb} {pkg.name} {pkg.version} (was: {pkg.state}, reasons: {', '.join(pkg.reasons)})"
            if dry_run:
                print(colors.dim(f"  [dry-run] would {label}"))
                log_line(log_file, f"[DRY-RUN] {verb} {pkg.name} {pkg.version} reasons={pkg.reasons}")
                continue
            try:
                run_mkp([verb] + args, mkp_bin)
                print(colors.green(f"  {verb} {pkg.name} {pkg.version}"))
                log_line(log_file, f"{verb} {pkg.name} {pkg.version} reasons={pkg.reasons}")
                if verb == "disable":
                    disabled += 1
                elif verb == "remove":
                    removed += 1
            except MkpError as exc:
                print(colors.red(f"  FAILED to {verb} {pkg.name} {pkg.version}: {exc}"))
                log_line(log_file, f"FAILED {verb} {pkg.name} {pkg.version}: {exc}")
                failed += 1
                break  # don't attempt remove if disable failed
    return disabled, removed, failed, skipped


def confirm(prompt: str) -> bool:
    try:
        answer = input(f"{prompt} [y/N]: ").strip().lower()
    except EOFError:
        return False
    return answer in ("y", "yes")


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog=PROG,
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  # just look around, no changes
  mkp-cleanup --list
  mkp-cleanup --view --filter '^my_plugin$'

  # dry-run a cleanup of packages that are enabled but inactive on this site
  mkp-cleanup --cleanup-enabled-inactive --dry-run

  # actually disable duplicate/outdated packages across all packages,
  # keeping the 2 newest versions of each, without confirmation prompt
  mkp-cleanup --cleanup-all-packages --keep-latest 2 --yes

  # disable AND permanently remove everything found in both areas
  mkp-cleanup --cleanup-all --purge --yes

  # use as a monitoring/cron gate (no changes, just exit code + one-liner)
  mkp-cleanup --list --nagios
""",
    )

    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--list", action="store_true", help="list outdated/duplicate MKP candidates (read-only)")
    mode.add_argument("--view", action="store_true", help="show detailed `mkp show` info for each candidate (read-only)")
    mode.add_argument("--cleanup-enabled-inactive", action="store_true",
                       help='clean up the "Enabled (inactive on this site)" area only')
    mode.add_argument("--cleanup-all-packages", action="store_true",
                       help='clean up the "All packages (enabled or disabled)" area '
                            '(inactive + superseded duplicates, plus until-version-exceeded '
                            'if --include-until-version-exceeded is set)')
    mode.add_argument("--cleanup-all", action="store_true",
                       help="clean up both areas in one run (union of the two switches above)")

    parser.add_argument("--filter", metavar="PATTERN", default=None,
                         help="only consider packages whose name matches this regex "
                              "(e.g. '^my_plugin$' for an exact match)")
    parser.add_argument("-i", "--ignore-case", action="store_true",
                         help="make --filter case-insensitive")
    parser.add_argument("--keep-latest", type=int, default=1, metavar="N",
                         help="number of newest versions to keep per package name "
                              "when detecting duplicates (default: 1)")
    parser.add_argument("--include-until-version-exceeded", action="store_true",
                         help="also treat packages whose 'Until Version' has been exceeded as "
                              "candidates (off by default: 'Until Version' is just a declaration "
                              "by the package author, the package may still work fine)")
    parser.add_argument("--dry-run", action="store_true",
                         help="show what would be done without changing anything")
    parser.add_argument("--purge", action="store_true",
                         help="also permanently `mkp remove` candidates after disabling them "
                              "(default: disable only, reversible via `mkp enable`)")
    parser.add_argument("-y", "--yes", action="store_true",
                         help="do not ask for confirmation before cleaning up")
    parser.add_argument("--json", action="store_true", help="machine-readable JSON output")
    parser.add_argument("--nagios", action="store_true",
                         help="Nagios/Checkmk-style OK/WARNING output + exit code; "
                             "only valid together with --list or --view")
    parser.add_argument("--log-file", metavar="PATH", default=None,
                         help=f"append actions to this log file (default: {default_log_path()})")
    parser.add_argument("--no-color", action="store_true", help="disable colored output")
    parser.add_argument("--skip-site-check", action="store_true",
                         help="skip the check that this is run inside a Checkmk site shell "
                              "(OMD_SITE/OMD_ROOT) - mainly for testing")
    parser.add_argument("--version", action="version", version=f"{PROG} {VERSION}")
    return parser


def main(argv=None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.nagios and not (args.list or args.view):
        parser.error("--nagios is only valid together with --list or --view (read-only modes)")
    if args.nagios and args.json:
        parser.error("--nagios and --json are mutually exclusive")
    if args.keep_latest < 1:
        parser.error("--keep-latest must be >= 1")

    colors = Colors(enabled=(not args.no_color) and sys.stdout.isatty() and not args.json and not args.nagios)
    log_file = args.log_file or default_log_path()

    pattern = None
    if args.filter:
        try:
            flags = re.IGNORECASE if args.ignore_case else 0
            pattern = re.compile(args.filter, flags)
        except re.error as exc:
            parser.error(f"invalid --filter regex: {exc}")

    try:
        check_site_context(args.skip_site_check)
        mkp_bin = find_mkp_binary()
        packages = fetch_packages(mkp_bin)
        current_version = get_current_cmk_version()
        classify(packages, args.keep_latest, current_version, args.include_until_version_exceeded)
    except MkpError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    enabled_inactive = select_enabled_inactive(packages, pattern)
    all_candidates = select_all_candidates(packages, pattern)

    # --- read-only modes ---------------------------------------------------
    if args.list or args.view:
        if args.nagios:
            return run_nagios(enabled_inactive, all_candidates)
        if args.json:
            return run_json(enabled_inactive, all_candidates, packages)
        run_list_or_view(args, enabled_inactive, all_candidates, mkp_bin, colors, current_version)
        return 0

    # --- cleanup modes -------------------------------------------------------
    areas = []
    if args.cleanup_enabled_inactive or args.cleanup_all:
        areas.append(("Enabled (inactive on this site)", enabled_inactive))
    if args.cleanup_all_packages or args.cleanup_all:
        areas.append(("All packages (enabled or disabled)", all_candidates))

    seen = set()
    to_clean = []
    for _, candidates in areas:
        for pkg in candidates:
            if pkg.key not in seen:
                seen.add(pkg.key)
                to_clean.append(pkg)

    print(colors.bold(f"{PROG}: found {len(to_clean)} candidate(s) across {len(areas)} area(s)"
                       f"{' matching filter ' + repr(args.filter) if args.filter else ''}"))
    for area_name, candidates in areas:
        print(f"\n{colors.bold(area_name)}:")
        print_table(["Name", "Version", "State", "Reason(s)"], candidate_rows(candidates))

    if not to_clean:
        print("\nnothing to do.")
        return 0

    action = "disable + remove (purge)" if args.purge else "disable"
    print(f"\nplanned action: {colors.yellow(action)}"
          f"{colors.dim(' [dry-run, no changes will be made]') if args.dry_run else ''}")

    already_disabled = sum(1 for p in to_clean if p.state == STATE_DISABLED)
    if not args.purge and already_disabled:
        noun = "package is" if already_disabled == 1 else "packages are"
        print(colors.yellow(
            f"note: {already_disabled} of {len(to_clean)} {noun} already Disabled - "
            f"'disable' is a no-op for those. Add --purge to actually remove them."
        ))

    if not args.dry_run and not args.yes:
        if not confirm(f"Proceed to {action} {len(to_clean)} package(s)?"):
            print("aborted.")
            return 1

    log_line(
        log_file,
        f"run: mode={'cleanup-all' if args.cleanup_all else ('cleanup-enabled-inactive' if args.cleanup_enabled_inactive else 'cleanup-all-packages')} "
        f"filter={args.filter!r} keep_latest={args.keep_latest} purge={args.purge} "
        f"include_until_version_exceeded={args.include_until_version_exceeded} "
        f"dry_run={args.dry_run} candidates={len(to_clean)}",
    )
    disabled, removed, failed, skipped = do_cleanup(to_clean, mkp_bin, args.dry_run, args.purge, log_file, colors)

    print(f"\n{colors.bold('summary')}: disabled={disabled} removed={removed} skipped={skipped} failed={failed}"
          f"{' (dry-run)' if args.dry_run else ''}")
    if skipped and not args.purge and not failed:
        print(colors.dim("hint: re-run with --purge to permanently remove the already-disabled duplicates."))
    return 0 if failed == 0 else 2


def run_list_or_view(args, enabled_inactive, all_candidates, mkp_bin, colors, current_version) -> None:
    print(colors.bold(f"{PROG} - current Checkmk version: {current_version or 'unknown'}"))
    if args.filter:
        print(colors.dim(f"filter: {args.filter!r} (case-{'insensitive' if args.ignore_case else 'sensitive'})"))

    print(f"\n{colors.bold('Enabled (inactive on this site)')} ({len(enabled_inactive)} found):")
    print_table(["Name", "Version", "State", "Reason(s)"], candidate_rows(enabled_inactive))

    print(f"\n{colors.bold('All packages (enabled or disabled)')} - outdated/duplicate candidates "
          f"({len(all_candidates)} found):")
    print_table(["Name", "Version", "State", "Reason(s)"], candidate_rows(all_candidates))

    if args.view:
        combined = {p.key: p for p in enabled_inactive + all_candidates}
        for pkg in sorted(combined.values(), key=lambda p: (p.name, version_key(p.version))):
            print(f"\n{colors.bold('---')} {pkg.name} {pkg.version} {colors.bold('---')}")
            print(f"  state:   {pkg.state}")
            print(f"  reasons: {', '.join(pkg.reasons)}")
            print(fetch_show(mkp_bin, pkg.name, pkg.version))


def run_json(enabled_inactive, all_candidates, packages) -> int:
    def dump(pkg: Package) -> dict:
        return {
            "name": pkg.name,
            "version": pkg.version,
            "title": pkg.title,
            "author": pkg.author,
            "req_version": pkg.req_version,
            "until_version": pkg.until_version,
            "files": pkg.files,
            "state": pkg.state,
            "reasons": pkg.reasons,
        }

    result = {
        "total_installed": len(packages),
        "enabled_inactive": [dump(p) for p in enabled_inactive],
        "all_candidates": [dump(p) for p in all_candidates],
    }
    print(json.dumps(result, indent=2))
    return 0


def run_nagios(enabled_inactive, all_candidates) -> int:
    combined = {p.key: p for p in enabled_inactive + all_candidates}
    count = len(combined)
    if count == 0:
        print("OK - no outdated/duplicate MKPs found")
        return 0
    names = ", ".join(f"{p.name} {p.version}" for p in list(combined.values())[:10])
    more = f" (+{count - 10} more)" if count > 10 else ""
    print(f"WARNING - {count} outdated/duplicate MKP(s) found: {names}{more} | mkp_candidates={count}")
    return 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\naborted.", file=sys.stderr)
        sys.exit(130)
