#!/usr/bin/env python3

import subprocess
import logging
import datetime
import os
import re
import sys
import shlex


# =========================
# Einstellungen
# =========================

ZFS_DATASETS = [
    "rpool/ROOT/pve-1",
    "rpool/pveconf"
]

SNAPSHOT_TAG = "pve-update-via-cron"
MAX_SNAPSHOTS = 5

LOGFILE = "/var/log/proxmox_update.log"

REBOOT_FLAG = os.environ.get("REBOOT", "NO").upper()

TIMESTAMP = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# Cron-Jobs laufen mit einem minimalen PATH (oft nur /usr/bin:/bin).
# proxmox-boot-tool & Co. liegen aber in /usr/sbin bzw. /sbin.
# Deshalb PATH hier explizit erweitern, damit es unter cron genauso
# funktioniert wie in einer interaktiven root-Shell.
os.environ["PATH"] = "/usr/sbin:/usr/bin:/sbin:/bin:" + os.environ.get("PATH", "")


# =========================
# Logging
# =========================

logging.basicConfig(
    filename=LOGFILE,
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)


def log(msg, level="info"):

    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    print(f"[{timestamp}] {msg}")

    if level == "info":
        logging.info(msg)

    elif level == "warn":
        logging.warning(msg)

    elif level == "error":
        logging.error(msg)


# =========================
# Kommando Helfer
# =========================

def run(cmd, check=True, quiet=False):
    """
    cmd darf entweder eine Liste von Argumenten sein (kein Shell-Aufruf,
    sicherer) oder ein String OHNE Shell-Metazeichen (Pipes, Redirects).
    Für Pipelines gibt es run_pipeline().

    quiet=True unterdrückt das automatische Loggen der stdout-Ausgabe
    (nützlich, wenn der Aufrufer die Ausgabe erst selbst filtert und
    danach gezielt loggt).
    """

    if isinstance(cmd, str):
        cmd = shlex.split(cmd)

    log(
        f"Ausführen: {' '.join(cmd)}"
    )

    try:

        result = subprocess.run(
            cmd,
            check=check,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if result.stdout.strip() and not quiet:
            log(result.stdout.strip())

        if result.stderr.strip() and not check:
            log(result.stderr.strip(), "warn")

        return result.stdout.strip()

    except FileNotFoundError as e:

        log(
            f"FEHLER: Befehl nicht gefunden: {' '.join(cmd)} ({e})",
            "error"
        )

        if check:
            raise

        return ""

    except subprocess.CalledProcessError as e:

        log(
            f"FEHLER: {' '.join(cmd)}",
            "error"
        )

        if e.stderr:
            log(e.stderr.strip(), "error")

        if check:
            raise

        return ""


def run_pipeline(cmd_str, check=False):
    """
    Führt einen Shell-Befehl MIT Pipes/Redirects tatsächlich über eine Shell
    aus (shell=True). Nur für fest im Skript definierte, ungefährliche
    Kommandos verwenden - niemals mit Benutzereingaben füttern.
    """

    log(f"Ausführen (shell): {cmd_str}")

    try:
        result = subprocess.run(
            cmd_str,
            shell=True,
            check=check,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if result.stdout.strip():
            log(result.stdout.strip())

        if result.stderr.strip():
            log(result.stderr.strip(), "warn")

        return result.stdout.strip()

    except subprocess.CalledProcessError as e:

        log(f"FEHLER: {cmd_str}", "error")

        if e.stderr:
            log(e.stderr.strip(), "error")

        if check:
            raise

        return ""


# =========================
# Kernel Prüfung
# =========================

def get_running_kernel():

    return run([
        "uname",
        "-r"
    ])


KERNEL_PKG_PATTERN = re.compile(r"^proxmox-kernel-([0-9][^\s]*?)(-signed)?$")


def _query_kernel_packages():
    """
    Fragt dpkg nach allen proxmox-kernel-* Paketen ab und gibt eine Liste
    von (status, paketname) Tupeln zurück. Die rohe dpkg-Ausgabe wird NICHT
    geloggt (quiet=True), da sie auch nicht-installierte ("un") und nur
    noch als Konfigurationsreste vorhandene ("rc") Pakete enthält, die wir
    im Log nicht als "installiert" anzeigen wollen.
    """

    output = run(
        ["dpkg-query", "-W", "-f=${db:Status-Abbrev} ${Package}\n",
         "proxmox-kernel-*"],
        check=False,
        quiet=True
    )

    packages = []

    for line in output.splitlines():

        parts = line.split(None, 1)

        if len(parts) != 2:
            continue

        status, pkg = parts

        packages.append((status, pkg))

    return packages


def get_installed_kernel():
    """
    Ermittelt die höchste installierte Kernel-Version rein in Python
    (kein fragiles Shell-Pipeline-Kommando mehr, kein sort -V nötig).
    Loggt dabei NUR die tatsächlich installierten ("ii") Kernel-Pakete,
    keine "un" (nicht installiert) oder "rc" (entfernt, Reste vorhanden).
    """

    packages = _query_kernel_packages()

    installed_pkgs = [pkg for status, pkg in packages if status == "ii"]

    if installed_pkgs:
        log(
            "Installierte Kernel-Pakete: " + ", ".join(sorted(installed_pkgs))
        )

    versions = []

    for pkg in installed_pkgs:

        m = KERNEL_PKG_PATTERN.match(pkg)

        if not m:
            continue

        versions.append(m.group(1))

    if not versions:
        return ""

    # Versionssortierung ohne externes sort -V:
    def version_key(v):
        return [int(p) if p.isdigit() else p for p in re.split(r"[.\-]", v)]

    versions.sort(key=version_key)

    return versions[-1]


def cleanup_old_kernel_configs():
    """
    Entfernt Konfigurationsreste von bereits deinstallierten alten
    Kernel-Paketen (dpkg-Status "rc" = removed, config-files still present).
    Diese Reste belegen unnötig Platz und tauchen sonst dauerhaft in
    'dpkg -l' auf, obwohl das Paket längst deinstalliert ist.
    """

    packages = _query_kernel_packages()

    residual_pkgs = [pkg for status, pkg in packages if status == "rc"]

    if not residual_pkgs:
        log(
            "OK: Keine Konfigurationsreste alter Kernel-Pakete gefunden"
        )
        return

    residual_pkgs = sorted(residual_pkgs)

    total = len(residual_pkgs)

    log(
        f"Entferne Konfigurationsreste von {total} alten Kernel-Paketen"
    )

    for index, pkg in enumerate(residual_pkgs, start=1):

        log(
            f"({index}/{total}) Lösche Reste von: {pkg}"
        )

        run(
            ["apt-get", "purge", "-y", "-qq", pkg],
            check=False,
            quiet=True
        )

        log(
            f"OK: ({index}/{total}) Erledigt: {pkg}"
        )


def reboot_required():

    running = get_running_kernel()

    installed = get_installed_kernel()

    log(
        f"Kernel aktiv: {running}"
    )

    log(
        f"Kernel installiert: {installed}"
    )

    if not running or not installed:
        log(
            "WARNUNG: Kernel-Vergleich nicht möglich (leerer Wert) - "
            "gehe sicherheitshalber von 'kein Reboot nötig' aus",
            "warn"
        )
        return False

    if running != installed:
        return True

    return False


# =========================
# Boot Prüfung
# =========================

def proxmox_boot_ok():

    try:
        result = subprocess.run(
            [
                "proxmox-boot-tool",
                "status"
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

    except FileNotFoundError as e:

        log(
            f"FEHLER: proxmox-boot-tool nicht gefunden (PATH-Problem?): {e}",
            "error"
        )

        return False

    if result.stdout:

        log(
            result.stdout.strip()
        )

    if result.returncode == 0:

        return True

    if result.stderr:

        log(
            result.stderr.strip(),
            "error"
        )
    else:
        log(
            f"FEHLER: proxmox-boot-tool status lieferte Exit-Code {result.returncode} ohne Fehlermeldung",
            "error"
        )

    return False


def refresh_boot():

    log(
        "Aktualisiere Proxmox Bootloader"
    )

    run([
        "proxmox-boot-tool",
        "refresh"
    ])


# =========================
# Update Prüfung
# =========================

def has_updates():

    run([
        "apt-get",
        "update"
    ])

    output = run(
        [
            "apt-get",
            "-s",
            "dist-upgrade"
        ],
        check=False
    )

    updates = [
        line
        for line in output.splitlines()
        if line.startswith("Inst ")
    ]

    if updates:

        log(
            f"{len(updates)} Updates verfügbar"
        )

        return True

    log(
        "OK: Keine Updates verfügbar"
    )

    return False


# =========================
# Snapshot Funktionen
# =========================

def create_snapshot(dataset):

    snapshot = (
        f"{dataset}@{SNAPSHOT_TAG}-{TIMESTAMP}"
    )

    log(
        f"Snapshot erstellen: {snapshot}"
    )

    run([
        "zfs",
        "snapshot",
        snapshot
    ])


def cleanup_snapshots(dataset):

    output = run(
        [
            "zfs",
            "list",
            "-H",
            "-t",
            "snapshot",
            "-o",
            "name",
            "-s",
            "creation"
        ],
        check=False
    )

    snapshots = []

    prefix = (
        f"{dataset}@{SNAPSHOT_TAG}-"
    )

    for line in output.splitlines():

        if line.startswith(prefix):

            snapshots.append(line)

    if len(snapshots) <= MAX_SNAPSHOTS:

        return

    for snap in snapshots[:-MAX_SNAPSHOTS]:

        log(
            f"Lösche Snapshot: {snap}"
        )

        run([
            "zfs",
            "destroy",
            "-r",
            snap
        ])


# =========================
# Upgrade
# =========================

def upgrade_system():

    log(
        "Starte dist-upgrade"
    )

    run([
        "apt-get",
        "dist-upgrade",
        "-y"
    ])

    log(
        "Starte autoremove"
    )

    run([
        "apt-get",
        "autoremove",
        "--purge",
        "-y"
    ])

    cleanup_old_kernel_configs()


# =========================
# Hauptprogramm
# =========================

def main():

    log(
        "===== Proxmox Update gestartet ====="
    )

    log(
        f"REBOOT-Flag: {REBOOT_FLAG}"
    )

    if os.geteuid() != 0:

        log(
            "FEHLER: Muss als root laufen",
            "error"
        )

        sys.exit(2)

    kernel_before = get_running_kernel()

    log(
        f"Kernel vor Update: {kernel_before}"
    )

    if has_updates():

        for dataset in ZFS_DATASETS:

            cleanup_snapshots(dataset)

            create_snapshot(dataset)

        upgrade_system()

        refresh_boot()

    else:

        log(
            "Keine Updates installiert"
        )

        # auch ohne Update-Lauf vorhandene Reste alter Kernel aufräumen
        cleanup_old_kernel_configs()

    # immer prüfen
    if reboot_required():

        log(
            "WARNUNG: Neuer Kernel aktivieren erforderlich",
            "warn"
        )

        if REBOOT_FLAG == "YES":

            log(
                "Prüfe proxmox-boot-tool"
            )

            if proxmox_boot_ok():

                log(
                    "OK: Bootloader OK"
                )

                log(
                    "REBOOT=YES erkannt"
                )

                run([
                    "shutdown",
                    "-r",
                    "+1",
                    "Proxmox-Update abgeschlossen"
                ])

            else:

                log(
                    "FEHLER: Bootloader Prüfung fehlgeschlagen - kein Neustart",
                    "error"
                )

        else:

            log(
                "REBOOT-Flag ist nicht 'YES' - kein automatischer Neustart"
            )

        sys.exit(1)

    log(
        "OK: Update erfolgreich abgeschlossen"
    )

    sys.exit(0)


if __name__ == "__main__":

    main()