"""Snapshots als Dateisystem zugaenglich machen.

Ein Proxmox-Snapshot ist kein Verzeichnis - je nach Storage muss er erst
zugaenglich gemacht werden:

    rbd (Ceph)   rbd map pool/image@snap        -> /dev/rbdN
    zfspool      zfs clone ds@snap tmp          -> /dev/zvol/tmp
                 (bei Container-Datasets direkt .zfs/snapshot/<snap>)
    lvmthin      lvchange -ay vg/snap_vol_snap  -> /dev/vg/snap_...
    dir/nfs      qemu-nbd --load-snapshot        -> /dev/nbdN

Danach werden die Partitionen ermittelt und **schreibgeschuetzt** gemountet.
Alles, was angelegt wurde, merkt sich diese Schicht und raeumt es wieder ab -
auch nach einem Absturz ueber `cleanup_leftovers()`.

Auf den Inhalt greifen Explorer und Web-Oberflaeche dann nur noch als
gewoehnlichen Pfad zu.
"""

from __future__ import annotations

import json
import logging
import os
import re
import shutil
import subprocess
import time
from dataclasses import dataclass, field

log = logging.getLogger("pvesnap.snapfs")

RUN_DIR = "/run/pvesnap"
MOUNT_ROOT = os.path.join(RUN_DIR, "mnt")
REGISTRY = os.path.join(RUN_DIR, "explorer.json")

# Datentraeger-Schluessel in den Gast-Konfigurationen
QEMU_DISK_KEY = re.compile(r"^(ide|sata|scsi|virtio)\d+$")
LXC_DISK_KEY = re.compile(r"^(rootfs|mp\d+)$")

# Mount-Optionen je Dateisystem. Snapshots laufender Gaeste haben fast immer
# ein unsauberes Journal - ohne diese Optionen wuerde der Kernel versuchen,
# es zurueckzuschreiben.
MOUNT_OPTIONS = {
    "ext2": ["ro"],
    "ext3": ["ro", "noload"],
    "ext4": ["ro", "noload"],
    "xfs": ["ro", "norecovery", "nouuid"],
    "btrfs": ["ro", "nologreplay"],
    "vfat": ["ro"],
    "ntfs": ["ro"],
    "ntfs3": ["ro"],
    "iso9660": ["ro"],
}
DEFAULT_MOUNT_OPTIONS = ["ro"]

MOUNTABLE = set(MOUNT_OPTIONS) | {"ext4", "xfs", "btrfs"}
IGNORED_FSTYPES = {"", "swap", "linux_raid_member", "crypto_LUKS", "LVM2_member"}


class SnapfsError(Exception):
    """Ein Snapshot liess sich nicht zugaenglich machen."""


# ---------------------------------------------------------------------------
# Kommandos
# ---------------------------------------------------------------------------

def run(args, timeout=120, check=True):
    """Externes Kommando ausfuehren; gibt die Standardausgabe zurueck."""
    log.debug("exec: %s", " ".join(args))
    try:
        proc = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                              timeout=timeout)
    except FileNotFoundError:
        raise SnapfsError("Kommando nicht gefunden: %s" % args[0])
    except subprocess.TimeoutExpired:
        raise SnapfsError("Zeitueberschreitung bei: %s" % " ".join(args))

    stdout = proc.stdout.decode("utf-8", "replace").strip()
    stderr = proc.stderr.decode("utf-8", "replace").strip()
    if check and proc.returncode != 0:
        raise SnapfsError("%s: %s" % (args[0], (stderr or stdout or "Fehler").splitlines()[0]))
    return stdout


def have(command):
    return shutil.which(command) is not None


# ---------------------------------------------------------------------------
# Datenmodell
# ---------------------------------------------------------------------------

@dataclass
class Volume:
    """Ein Datentraeger aus der Gast-Konfiguration."""
    key: str            # scsi0, rootfs, mp0 ...
    volid: str          # data:vm-802-disk-0
    storage: str        # data
    volname: str        # vm-802-disk-0
    size: str = ""

    @property
    def label(self):
        return "%s  %s%s" % (self.key.ljust(8), self.volid,
                             ("  (%s)" % self.size) if self.size else "")


@dataclass
class Device:
    """Ein Blockgeraet (ganze Platte oder Partition) des Snapshots."""
    path: str
    name: str = ""
    size: int = 0
    fstype: str = ""
    label: str = ""
    kind: str = "part"
    volume: object = None

    @property
    def mountable(self):
        return bool(self.fstype) and self.fstype not in IGNORED_FSTYPES

    @property
    def human_size(self):
        return human_bytes(self.size)

    @property
    def title(self):
        parts = [os.path.basename(self.path)]
        if self.label:
            parts.append('"%s"' % self.label)
        parts.append(self.fstype or "unbekannt")
        parts.append(self.human_size)
        return "  ".join(parts)


def human_bytes(count):
    value = float(count or 0)
    for unit in ("B", "K", "M", "G", "T", "P"):
        if value < 1024 or unit == "P":
            return "%d%s" % (value, unit) if unit == "B" else "%.1f%s" % (value, unit)
        value /= 1024
    return "%.1fP" % value


# ---------------------------------------------------------------------------
# Konfiguration des Snapshots lesen
# ---------------------------------------------------------------------------

def snapshot_config(proxmox, guest, snapname):
    path = "/nodes/%s/%s/%d/snapshot/%s/config" % (guest.node, guest.type,
                                                   guest.vmid, snapname)
    data = proxmox._json(["get", path])
    if not isinstance(data, dict):
        raise SnapfsError("Konfiguration des Snapshots %r nicht lesbar" % snapname)
    return data


def volumes_from_config(config, guest_type):
    """Alle echten Datentraeger aus einer Gast-Konfiguration."""
    pattern = LXC_DISK_KEY if guest_type == "lxc" else QEMU_DISK_KEY
    volumes = []
    for key in sorted(config):
        if not pattern.match(key):
            continue
        value = str(config[key])
        if not value or "media=cdrom" in value:
            continue
        volid = value.split(",", 1)[0].strip()
        if volid in ("none", "cdrom") or ":" not in volid:
            continue
        storage, volname = volid.split(":", 1)
        size = ""
        for part in value.split(",")[1:]:
            if part.strip().startswith("size="):
                size = part.strip()[5:]
        volumes.append(Volume(key=key, volid=volid, storage=storage,
                              volname=volname, size=size))
    return volumes


def storage_info(proxmox, name):
    data = proxmox._json(["get", "/storage/%s" % name])
    if not isinstance(data, dict):
        raise SnapfsError("Storage %r nicht gefunden" % name)
    return data


# ---------------------------------------------------------------------------
# Merkliste, damit nach einem Absturz nichts stehen bleibt
# ---------------------------------------------------------------------------

def _registry_load():
    try:
        with open(REGISTRY, "r", encoding="utf-8") as handle:
            data = json.load(handle)
        return data if isinstance(data, list) else []
    except (OSError, ValueError):
        return []


def _registry_save(entries):
    try:
        os.makedirs(RUN_DIR, exist_ok=True)
        tmp = REGISTRY + ".tmp"
        with open(tmp, "w", encoding="utf-8") as handle:
            json.dump(entries, handle, indent=1)
        os.replace(tmp, REGISTRY)
    except OSError as exc:
        log.debug("Merkliste nicht schreibbar: %s", exc)


def _registry_add(kind, value, extra=None):
    entries = _registry_load()
    entries.append({"kind": kind, "value": value, "extra": extra or {},
                    "pid": os.getpid(), "time": int(time.time())})
    _registry_save(entries)


def _registry_remove(kind, value):
    entries = [e for e in _registry_load()
               if not (e.get("kind") == kind and e.get("value") == value)]
    _registry_save(entries)


def cleanup_leftovers(verbose=False):
    """Raeumt auf, was ein abgestuerzter Lauf hinterlassen hat."""
    removed = []
    for entry in reversed(_registry_load()):
        kind, value = entry.get("kind"), entry.get("value")
        extra = entry.get("extra") or {}
        try:
            if kind == "mount":
                _umount(value)
            elif kind == "rbd":
                run(["rbd", "unmap", value], check=False)
            elif kind == "rbd-nbd":
                run(["rbd-nbd", "unmap", value], check=False)
            elif kind == "nbd":
                run(["qemu-nbd", "--disconnect", value], check=False)
            elif kind == "lvm":
                run(["lvchange", "-an", value], check=False)
            elif kind == "zfsclone":
                run(["zfs", "destroy", value], check=False)
            elif kind == "dir":
                if os.path.isdir(value):
                    os.rmdir(value)
            else:
                continue
            removed.append("%s %s" % (kind, value))
            if verbose:
                print("aufgeraeumt: %s %s" % (kind, value))
        except (SnapfsError, OSError) as exc:
            if verbose:
                print("konnte %s %s nicht aufraeumen: %s" % (kind, value, exc))
        _registry_remove(kind, value)
        del extra
    return removed


def _umount(mountpoint):
    if not os.path.ismount(mountpoint):
        return
    try:
        run(["umount", mountpoint])
    except SnapfsError:
        run(["umount", "-l", mountpoint], check=False)   # notfalls verzoegert


# ---------------------------------------------------------------------------
# Sitzung
# ---------------------------------------------------------------------------

class Session:
    """Ein geoeffneter Snapshot: eingebundene Geraete und Mountpunkte."""

    def __init__(self, proxmox, guest, snapname):
        self.proxmox = proxmox
        self.guest = guest
        self.snapname = snapname
        self.volumes = []
        self.devices = []
        self.mounts = {}            # Geraetepfad -> Mountpunkt
        self.notes = []             # Hinweise fuer die Oberflaeche
        self._cleanup = []          # (art, wert) in umgekehrter Reihenfolge

    # -- oeffnen ----------------------------------------------------------

    def open(self):
        config = snapshot_config(self.proxmox, self.guest, self.snapname)
        self.volumes = volumes_from_config(config, self.guest.type)
        if not self.volumes:
            raise SnapfsError("Der Snapshot enthaelt keine Datentraeger.")
        return self.volumes

    def prepare(self):
        """Alle Datentraeger einbinden und mountbare Dateisysteme mounten."""
        for volume in self.volumes:
            try:
                self._prepare_volume(volume)
            except SnapfsError as exc:
                self.notes.append("%s: %s" % (volume.key, exc))
                log.warning("%s (%s): %s", volume.key, volume.volid, exc)
        if not self.mounts and not self.devices:
            raise SnapfsError("Kein Datentraeger des Snapshots konnte "
                              "eingebunden werden.\n- " + "\n- ".join(self.notes))
        return self.mounts

    def _prepare_volume(self, volume):
        info = storage_info(self.proxmox, volume.storage)
        kind = info.get("type")

        if kind == "zfspool" and self.guest.type == "lxc":
            # Container liegen auf ZFS als Dateisystem - der Snapshot ist ohne
            # Umweg unter .zfs/snapshot/<name> sichtbar.
            path = self._zfs_snapshot_dir(info, volume)
            if path:
                self.mounts[path] = path
                self.notes.append("%s: direkt aus %s" % (volume.key, path))
                return

        device_path = self._map(info, volume)
        for device in self._scan(device_path, volume):
            self.devices.append(device)
            if device.mountable:
                try:
                    self.mount(device)
                except SnapfsError as exc:
                    self.notes.append("%s: %s" % (device.path, exc))

    # -- Einbinden je Storage-Typ -----------------------------------------

    def _map(self, info, volume):
        kind = info.get("type")
        if kind == "rbd":
            return self._map_rbd(info, volume)
        if kind == "zfspool":
            return self._map_zfs(info, volume)
        if kind in ("lvmthin", "lvm"):
            return self._map_lvm(info, volume)
        if kind in ("dir", "nfs", "cifs", "glusterfs"):
            return self._map_qcow(info, volume)
        raise SnapfsError("Storage-Typ %r wird vom Explorer nicht unterstuetzt" % kind)

    def _rbd_args(self, info):
        args = []
        storage = info.get("storage") or info.get("name") or ""
        keyring = "/etc/pve/priv/ceph/%s.keyring" % storage
        if info.get("monhost"):
            args += ["-m", str(info["monhost"]).replace(" ", ",")]
            args += ["--id", info.get("username") or "admin"]
            if os.path.exists(keyring):
                args += ["--keyring", keyring]
        else:
            args += ["--id", info.get("username") or "admin"]
        return args

    def _map_rbd(self, info, volume):
        if not have("rbd"):
            raise SnapfsError("'rbd' ist nicht installiert")
        pool = info.get("pool") or "rbd"
        namespace = info.get("namespace")
        image = "%s/%s%s@%s" % (pool, (namespace + "/") if namespace else "",
                                volume.volname, self.snapname)
        base = self._rbd_args(info)

        try:
            output = run(["rbd"] + base + ["map", "--read-only", image])
            device = self._first_device(output)
            self._remember("rbd", device)
            return device
        except SnapfsError as exc:
            # krbd kann manche Image-Features nicht - rbd-nbd kann alle.
            if not have("rbd-nbd"):
                raise SnapfsError("%s (rbd-nbd waere die Alternative, ist aber "
                                  "nicht installiert: apt install ceph-common)" % exc)
            log.info("rbd map fehlgeschlagen (%s) - versuche rbd-nbd", exc)
            output = run(["rbd-nbd"] + base + ["map", "--read-only", image], timeout=180)
            device = self._first_device(output)
            self._remember("rbd-nbd", device)
            return device

    def _zfs_dataset(self, info, volume):
        pool = info.get("pool") or info.get("zfspool") or ""
        return "%s/%s" % (pool.rstrip("/"), volume.volname)

    def _zfs_snapshot_dir(self, info, volume):
        dataset = self._zfs_dataset(info, volume)
        try:
            mountpoint = run(["zfs", "get", "-H", "-o", "value", "mountpoint", dataset])
        except SnapfsError:
            return ""
        if not mountpoint or mountpoint in ("none", "-", "legacy"):
            return ""
        candidate = os.path.join(mountpoint, ".zfs", "snapshot", self.snapname)
        return candidate if os.path.isdir(candidate) else ""

    def _map_zfs(self, info, volume):
        if not have("zfs"):
            raise SnapfsError("'zfs' ist nicht installiert")
        dataset = self._zfs_dataset(info, volume)
        source = "%s@%s" % (dataset, self.snapname)
        clone = "%s-pvesnap-%d" % (dataset, os.getpid())
        run(["zfs", "clone", "-o", "readonly=on", source, clone])
        self._remember("zfsclone", clone)
        device = "/dev/zvol/%s" % clone
        for _ in range(50):                       # udev braucht einen Moment
            if os.path.exists(device):
                return device
            time.sleep(0.1)
        raise SnapfsError("Klon %s ist nicht als Geraet aufgetaucht" % clone)

    def _map_lvm(self, info, volume):
        if not have("lvchange"):
            raise SnapfsError("LVM-Werkzeuge sind nicht installiert")
        vg = info.get("vgname") or ""
        # Proxmox benennt Snapshot-LVs immer nach diesem Schema.
        lv = "snap_%s_%s" % (volume.volname, self.snapname)
        target = "%s/%s" % (vg, lv)
        run(["lvchange", "-ay", "-K", "--readonly", target])
        self._remember("lvm", target)
        device = "/dev/%s/%s" % (vg, lv)
        if not os.path.exists(device):
            raise SnapfsError("%s ist nach dem Aktivieren nicht vorhanden" % device)
        return device

    def _map_qcow(self, info, volume):
        if not have("qemu-nbd"):
            raise SnapfsError("'qemu-nbd' ist nicht installiert (apt install qemu-utils)")
        base = info.get("path") or ""
        image = os.path.join(base, "images", str(self.guest.vmid), volume.volname)
        if not os.path.exists(image):
            raise SnapfsError("Abbilddatei nicht gefunden: %s" % image)
        if not image.endswith(".qcow2"):
            raise SnapfsError("nur qcow2 kann Snapshots enthalten (%s)"
                              % os.path.basename(image))

        run(["modprobe", "nbd", "max_part=16"], check=False)
        device = self._free_nbd()
        run(["qemu-nbd", "--read-only", "--load-snapshot=%s" % self.snapname,
             "--connect=%s" % device, image], timeout=180)
        self._remember("nbd", device)
        time.sleep(0.5)
        return device

    @staticmethod
    def _free_nbd():
        for index in range(16):
            device = "/dev/nbd%d" % index
            if not os.path.exists(device):
                continue
            try:
                with open("/sys/block/nbd%d/size" % index) as handle:
                    if handle.read().strip() == "0":
                        return device
            except OSError:
                continue
        raise SnapfsError("kein freies /dev/nbdN vorhanden")

    @staticmethod
    def _first_device(output):
        for line in (output or "").splitlines():
            line = line.strip()
            if line.startswith("/dev/"):
                return line
        raise SnapfsError("Antwort ohne Geraetenamen: %s" % (output or "(leer)"))

    # -- Partitionen und Dateisysteme -------------------------------------

    def _scan(self, device_path, volume):
        """Partitionen des eingebundenen Geraets ermitteln."""
        devices = self._lsblk(device_path, volume)
        if len(devices) == 1 and not devices[0].mountable:
            # Partitionstabelle vorhanden, aber noch nicht eingelesen
            for tool in (["partx", "-a", device_path], ["kpartx", "-a", "-r", device_path]):
                if have(tool[0]):
                    run(tool, check=False)
                    time.sleep(0.5)
                    devices = self._lsblk(device_path, volume)
                    if len(devices) > 1:
                        break
        return devices

    def _lsblk(self, device_path, volume):
        try:
            output = run(["lsblk", "-J", "-b", "-o",
                          "NAME,PATH,SIZE,FSTYPE,LABEL,TYPE", device_path])
            data = json.loads(output)
        except (SnapfsError, ValueError) as exc:
            raise SnapfsError("Partitionen nicht lesbar: %s" % exc)

        found = []

        def walk(entries):
            for entry in entries:
                device = Device(
                    path=entry.get("path") or "",
                    name=entry.get("name") or "",
                    size=int(entry.get("size") or 0),
                    fstype=entry.get("fstype") or "",
                    label=entry.get("label") or "",
                    kind=entry.get("type") or "part",
                    volume=volume,
                )
                children = entry.get("children") or []
                if device.path:
                    # Eine Platte mit Partitionen selbst nicht anbieten
                    if not (children and not device.fstype):
                        found.append(device)
                if device.fstype == "LVM2_member":
                    self.notes.append(
                        "%s enthaelt LVM - im Gast angelegte Volume-Groups werden "
                        "nicht automatisch aktiviert (Namenskonflikte mit dem Host)"
                        % device.path)
                walk(children)

        walk(data.get("blockdevices") or [])
        return found

    def mount(self, device):
        if device.path in self.mounts:
            return self.mounts[device.path]

        name = "%s-%d-%s-%s" % (self.guest.type, self.guest.vmid,
                                re.sub(r"[^A-Za-z0-9_.-]", "_", self.snapname),
                                os.path.basename(device.path))
        mountpoint = os.path.join(MOUNT_ROOT, name)
        os.makedirs(mountpoint, exist_ok=True)
        self._remember("dir", mountpoint)

        options = MOUNT_OPTIONS.get(device.fstype, DEFAULT_MOUNT_OPTIONS)
        attempts = [options]
        if options != DEFAULT_MOUNT_OPTIONS:
            attempts.append(DEFAULT_MOUNT_OPTIONS)   # notfalls schlicht "ro"

        last = None
        for option_set in attempts:
            try:
                run(["mount", "-o", ",".join(option_set), "-t", device.fstype or "auto",
                     device.path, mountpoint])
                self.mounts[device.path] = mountpoint
                self._remember("mount", mountpoint)
                log.info("%s (%s) eingehaengt unter %s",
                         device.path, device.fstype, mountpoint)
                return mountpoint
            except SnapfsError as exc:
                last = exc
        raise SnapfsError("%s nicht einhaengbar: %s" % (device.path, last))

    # -- aufraeumen -------------------------------------------------------

    def _remember(self, kind, value):
        self._cleanup.append((kind, value))
        _registry_add(kind, value)

    def close(self):
        for kind, value in reversed(self._cleanup):
            try:
                if kind == "mount":
                    _umount(value)
                elif kind == "rbd":
                    run(["rbd", "unmap", value], check=False)
                elif kind == "rbd-nbd":
                    run(["rbd-nbd", "unmap", value], check=False)
                elif kind == "nbd":
                    run(["qemu-nbd", "--disconnect", value], check=False)
                elif kind == "lvm":
                    run(["lvchange", "-an", value], check=False)
                elif kind == "zfsclone":
                    run(["zfs", "destroy", value], check=False)
                elif kind == "dir" and os.path.isdir(value):
                    os.rmdir(value)
            except (SnapfsError, OSError) as exc:
                log.warning("Aufraeumen von %s %s fehlgeschlagen: %s", kind, value, exc)
            _registry_remove(kind, value)
        self._cleanup = []
        self.mounts = {}

    def __enter__(self):
        return self

    def __exit__(self, *_exc):
        self.close()
        return False


def list_snapshots(proxmox, guest):
    """Snapshots eines Gastes, neueste zuerst."""
    snapshots = proxmox.list_snapshots(guest)
    return sorted(snapshots, key=lambda s: s.snaptime, reverse=True)