#!/bin/bash
#
# shellcheck disable=SC2181
#
# Copyright (c) 2024-2026 openmediavault plugin developers
#
# This file is licensed under the terms of the GNU General Public
# License version 2. This program is licensed "as is" without any
# warranty of any kind, whether express or implied.
#
# version: 0.1.0
#

if [[ $(id -u) -ne 0 ]]; then
  echo "This script must be executed as root or using sudo."
  exit 99
fi

export LC_ALL=C.UTF-8

# declare/initialize
declare -i poweredOn=0

action=""
cntrName=""
snapDir=""
snapName=""
virshCmd="virsh --connect lxc:///"

# logging location
logDir="/var/log/"
logFile="${logDir}/omv-lxc-snapshot.log"

_log()
{
  local msg=${1}
  echo "[$(date +'%Y-%m-%d %H:%M:%S%z')] [lxcsnapshot] ${msg}" | tee -a "${logFile}" >&2
}

# Filesystem helpers
_fs_type()
{
  findmnt -no FSTYPE --target "${1}" 2>/dev/null
}

_mount_uuid()
{
  findmnt -no UUID --target "${1}" 2>/dev/null
}

_is_btrfs_fs()
{
  [[ "$(_fs_type "${1}")" == "btrfs" ]]
}

_is_btrfs_subvol()
{
  btrfs subvolume show "${1}" >/dev/null 2>&1
}

_same_fs()
{
  # same filesystem (UUID) requirement for btrfs snapshot source/target
  local a_uuid b_uuid
  a_uuid="$(_mount_uuid "${1}")"
  b_uuid="$(_mount_uuid "${2}")"
  [[ -n "${a_uuid}" && -n "${b_uuid}" && "${a_uuid}" == "${b_uuid}" ]]
}

_btrfs_revert_snapshot()
{
  local src_snap="${1}" dest_subvol="${2}"
  # Replace dest_subvol with a RW snapshot of src_snap
  # Ensure dest is currently a subvolume (or remove dir and create)
  if _is_btrfs_subvol "${dest_subvol}"; then
    btrfs subvolume delete "${dest_subvol}" || return 1
  else
    rm -rf "${dest_subvol}" || return 1
  fi
  btrfs subvolume snapshot "${src_snap}" "${dest_subvol}"
}

# loop through options
while getopts "a:c:d:hs:" opt; do
  case "${opt}" in
    a)
      action="${OPTARG}"
      ;;
    c)
      cntrName="${OPTARG}"
      ;;
    d)
      snapDir="${OPTARG}"
      ;;
    h)
      echo "Use the following flags:"
      echo "  -a action"
      echo "    action can be 'create', 'revert', 'list', or 'delete'"
      echo "  -c name"
      echo "    container name"
      echo "  -d 'dir'"
      echo "    set snapshot directory (must be btrfs subvolume to use btrfs snapshots)."
      echo "    a subdirectory named after container name will be created if it does not exist."
      echo "    if a directory is not specified, the container's parent directory will be used"
      echo "    and a snapshots subdirectory will be created."
      echo "  -h"
      echo "    show this help"
      echo "  -s 'snapshot name'"
      echo ""
      echo "Examples:"
      echo "  omv-lxc-snapshot -c testlxc1 -a create -d '/srv/lxc/snapshots/' -s '20231123100102'"
      echo ""
      exit 100
      ;;
    s)
      snapName="${OPTARG}"
      ;;
    \?)
      echo "Invalid option: -${OPTARG}"
      ;;
  esac
done

if [ ! -f "/usr/bin/virsh" ]; then
  _log "virsh command not found!"
  exit 2
fi

if [ ! -f "/usr/bin/btrfs" ]; then
  _log "btrfs command not found!"
  exit 3
fi

# check if container name was specified
if [ -z "${cntrName}" ]; then
  _log "No container specified!"
  exit 4
fi

# check if VM is defined by libvirt
if ! ${virshCmd} dominfo "${cntrName}" > /dev/null 2>&1; then
  _log "VM '${cntrName}' not defined!"
  exit 5
fi

cntrDir="$(${virshCmd} dumpxml "${cntrName}" | xmlstarlet sel -t -m "/domain/devices/filesystem[target/@dir='/']" -v "source/@dir" -n)"

# if container path doesn't exist, exit
if [ "${action}" = "create" ] && [ ! -d "${cntrDir}" ]; then
  _log "${cntrDir} does not exist.  Exiting."
  exit 6
fi

# choose snapshot directory default
if [ -z "${snapDir}" ]; then
  parentDir="$(dirname "${cntrDir}")"
  if [ "${parentDir}" = "/" ]; then
    parentDir=""
  fi
  snapDir="${parentDir}/snapshots"
fi

# ensure snapshot dir for container exists
cntrSnapDir="${snapDir}/${cntrName}"
mkdir -p "${cntrSnapDir}"

# Decide snapshot mode
use_btrfs="no"
snap_mode="tar"
if _is_btrfs_fs "${cntrDir}" && _is_btrfs_subvol "${cntrDir}" && _same_fs "${cntrDir}" "${cntrSnapDir}"; then
  use_btrfs="yes"
  snap_mode="btrfs snapshot"
fi

_is_running()
{
  ${virshCmd} list --name | grep -qx "${cntrName}"
}

# Helper to stop/start container if running
# returns non-zero if the container did not stop
_stop_if_running()
{
  local countDown=60
  if _is_running; then
    _log "Shutdown ${cntrName} ..."
    ${virshCmd} shutdown "${cntrName}"
    poweredOn=1
    # wait for the container to actually stop before touching its rootfs
    while _is_running; do
      if [ ${countDown} -le 0 ]; then
        _log "${cntrName} did not shut down within 60 seconds."
        return 1
      fi
      sleep 2
      countDown=$((countDown-2))
    done
  fi
  return 0
}

_start_if_needed()
{
  if [ ${poweredOn} -eq 1 ]; then
    _log "Starting ${cntrName} ..."
    ${virshCmd} start "${cntrName}"
  fi
}

case "${action}" in
  create)
    # create snapshot name if none specified
    if [ -z "${snapName}" ]; then
      snapName="$(date +'%Y%m%d_%H%M%S')"
    fi

    _log "Starting snapshot '${snapName}' for container '${cntrName}' (mode: ${snap_mode})"
    if ! _stop_if_running; then
      _log "Aborting snapshot."
      exit 17
    fi

    if [ "${use_btrfs}" = "yes" ]; then
      # store btrfs snapshots in a subdir to avoid mixing files and subvols
      # e.g. <snapDir>/<container>/<snapshotName>
      dst="${cntrSnapDir}/${snapName}"
      _log "Creating Btrfs readonly snapshot at ${dst} ..."
      if ! btrfs subvolume snapshot -r "${cntrDir}" "${dst}"; then
        _log "Btrfs snapshot failed; falling back to tar.gz"
        use_btrfs="no"
        snap_mode="tar"
      fi
    fi

    if [ "${use_btrfs}" != "yes" ]; then
      snapFile="${cntrSnapDir}/${snapName}.tar.gz"
      _log "Creating ${snapFile} ..."
      tar czf "${snapFile}" -C "${cntrDir}" .
      if [ $? -gt 0 ]; then
        _log "creating archive failed."
        exit 10
      fi
    fi

    _start_if_needed
    _log "Snapshot '${snapName}' complete."
    ;;

  revert)
    if [ -z "${snapName}" ]; then
      _log "No snapshot name specified!"
      exit 11
    fi
    _log "Reverting to snapshot '${snapName}' for container '${cntrName}'"

    if ! _stop_if_running; then
      _log "Aborting revert; container files were not changed."
      exit 17
    fi

    if [ "${use_btrfs}" = "yes" ] && [ -d "${cntrSnapDir}/${snapName}" ]; then
      # Btrfs subvolume snapshot exists
      src="${cntrSnapDir}/${snapName}"
      _log "Reverting using Btrfs snapshot ${src} -> ${cntrDir} ..."
      if ! _btrfs_revert_snapshot "${src}" "${cntrDir}"; then
        _log "Btrfs revert failed."
        exit 12
      fi
    else
      # tar.gz fallback
      snapFile="${cntrSnapDir}/${snapName}.tar.gz"
      if [ ! -f "${snapFile}" ]; then
        _log "Snapshot not found: ${snapFile}"
        exit 13
      fi
      _log "Removing files in ${cntrDir} ..."
      rm -rf "${cntrDir}" 2>/dev/null
      sleep 2
      _log "Reverting from ${snapFile} ..."
      mkdir -p "${cntrDir}"
      tar xzf "${snapFile}" -C "${cntrDir}"
      if [ $? -gt 0 ]; then
        _log "extracting archive failed."
        exit 14
      fi
    fi

    _start_if_needed
    _log "Snapshot '${snapName}' revert complete."
    ;;

  list)
    # show btrfs and tar snapshots distinctly
    if [ -d "${cntrSnapDir}" ]; then
      # Btrfs snapshots = subdirectories that are subvolumes
      while IFS= read -r d; do
        _is_btrfs_subvol "${d}" && echo "$(basename "${d}") (btrfs)"
      done < <(find "${cntrSnapDir}" -mindepth 1 -maxdepth 1 -type d | sort)
      # tar snapshots
      find "${cntrSnapDir}" -type f -name "*.tar.gz" -printf "%f\n" | sed 's/.tar.gz$//' | sed 's/$/ (tar)/' | sort
    fi
    ;;

  delete)
    if [ -z "${snapName}" ]; then
      _log "No snapshot name specified!"
      exit 15
    fi
    _log "Deleting snapshot '${snapName}' for container '${cntrName}'"
    if [ -d "${cntrSnapDir}/${snapName}" ] && _is_btrfs_subvol "${cntrSnapDir}/${snapName}"; then
      _log "Deleting Btrfs snapshot ${cntrSnapDir}/${snapName} ..."
      if ! btrfs subvolume delete "${cntrSnapDir}/${snapName}"; then
        _log "Failed to delete btrfs snapshot."
        exit 16
      fi
    else
      _log "Deleting tar snapshot ${cntrSnapDir}/${snapName}.tar.gz ..."
      rm -fv "${cntrSnapDir}/${snapName}.tar.gz"
    fi

    _log "Snapshot '${snapName}' delete complete."
    ;;

  *)
    echo "Error: Invalid action. Should be 'create', 'revert', 'list', or 'delete'."
    exit 1
    ;;
esac

exit 0
