#!/usr/bin/env python3
import subprocess
import logging
import datetime
import os
import sys
import re

# ========== Einstellungen ==========
ZFS_DATASETS = ["rpool/ROOT", "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")

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

def run(cmd, check=True):
    """Kommando ausführen"""
    logging.debug(f"Running: {cmd}")
    try:
        result = subprocess.run(
            cmd, shell=True, check=check,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        return result.stdout.decode().strip()
    except subprocess.CalledProcessError as e:
        logging.error(f"Fehler bei Befehl: {cmd}")
        logging.error(e.stderr.decode().strip())
        if check:
            raise
        return ""

def has_updates():
    """apt update ausführen & prüfen, ob Updates vorhanden"""
    try:
        output = run("apt update")
        if "upgradable" in output.lower():
            logging.info("Updates verfügbar.")
            return True
        logging.info("Keine Updates verfügbar.")
        return False
    except Exception:
        logging.error("apt update fehlgeschlagen.")
        sys.exit(2)

def is_kernel_update_available():
    """Prüfen, ob ein Kernel-Update verfügbar ist"""
    output = run("apt list --upgradable", check=False)
    return any("linux-image-" in line for line in output.splitlines())

def create_snapshot(dataset):
    """Snapshot mit Tag erstellen"""
    snapshot_name = f"{dataset}@{SNAPSHOT_TAG}-{TIMESTAMP}"
    logging.info(f"Erstelle Snapshot: {snapshot_name}")
    run(f"zfs snapshot {snapshot_name}")

def cleanup_snapshots(dataset):
    """Alte Snapshots mit dem definierten Tag löschen"""
    logging.info(f"Bereinige alte Snapshots in {dataset}")
    snap_list = run(f"zfs list -t snapshot -o name -s creation | grep '^{dataset}@{SNAPSHOT_TAG}-'", check=False).splitlines()
    if len(snap_list) > MAX_SNAPSHOTS:
        for snap in snap_list[:-MAX_SNAPSHOTS]:
            logging.info(f"Lösche Snapshot: {snap}")
            run(f"zfs destroy {snap}")

def upgrade_system():
    """Upgrade durchführen"""
    try:
        logging.info("Starte apt dist-upgrade …")
        run("apt dist-upgrade -y")
        logging.info("Führe apt autoremove aus …")
        run("apt autoremove -y")
    except Exception:
        logging.error("Upgrade fehlgeschlagen.")
        sys.exit(2)

def main():
    logging.info("===== Proxmox Update & Snapshot gestartet =====")

    if os.geteuid() != 0:
        print("Dieses Skript muss als root ausgeführt werden.")
        sys.exit(2)

    if not has_updates():
        logging.info("Beende – keine Updates.")
        sys.exit(0)

    # Snapshots nur bei Update
    for dataset in ZFS_DATASETS:
        create_snapshot(dataset)
        cleanup_snapshots(dataset)

    # Upgrade
    upgrade_system()

    # Kernel-Update erkennen
    if is_kernel_update_available():
        logging.warning("Kernel-Update erkannt – Neustart empfohlen.")
        if REBOOT_FLAG == "YES":
            logging.info("Reboot wird ausgeführt …")
            run("reboot")
        sys.exit(1)  # Exit 1 = Neustart empfohlen

    logging.info("===== Update-Vorgang abgeschlossen =====")
    sys.exit(0)

if __name__ == "__main__":
    main()