#!/bin/bash

# On Ian's computers, mount received subvolumes after btrbk.
# Copyright (C) 2024  Ian Kelling

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# SPDX-License-Identifier: GPL-3.0-or-later

this_file="$(readlink -f -- "${BASH_SOURCE[0]}")"
readonly this_file
cd /
[[ $EUID == 0 ]] || exec sudo -E "$this_file" "$@"

set -e; . /usr/local/lib/bash-bear; set +e

usage() {
  cat <<EOF
Usage: ${0##*/} [OPTIONS] [SUBVOLUMES]

-h|--help  Print help and exit.
-f|--force  Use kill -9 to try fixing unmount errors
-v|--verbose  Be more verbose


Note: In git this is not not executable because it's meant to be installed
using ./install-my-scripts

Note: Uses util-linux getopt option parsing: spaces between args and
options, short options can be combined, options before args.
EOF
  exit $1
}

pre="mount-latest-subvol:${SSH_CLIENT:+ $HOSTNAME:}"

tu() {
  while read -r line; do
    file="$1"
    grep -xFq "$line" "$file" || tee -a "$file"<<<"$line"
  done
}
d() {
  if $verbose; then
    printf "$pre %s\n" "$*"
  fi
}
m() {
  if $verbose; then
    printf "$pre %s\n" "$*"
  fi
  "$@"
}
x() {
  printf "$pre %s\n" "$*"
  "$@"
}

mnt() {
  dir=$1
  if ! mountpoint -q $dir; then
    mkdir -p $dir
    m mount $dir
  fi
}
fstab() {
  while read -r start mpoint end; do
    l="$start $mpoint $end"
    # kill off any lines that duplicate the mount point.
    sed --follow-symlinks -ri "\%$l%b;\%^\s*\S+\s+$mpoint\s%d" /etc/fstab
    tu /etc/fstab <<<"$l"
  done
}
pid-check() {
  for p in ${pids}; do
    for m in ${my_pids[@]}; do
      if (( p == m )); then
        echo "$0: error: pids to kill includes our pid or a parent. ps output:" >&2
        ps -f -p $p
        exit 1
      fi
    done
  done
}
kill-dir() {
  for sig; do
    echo kill-dir $sig
    found_pids=false
    if pids=$(timeout 4 lsof -t $dir); then
      found_pids=true
      timeout 4 lsof -w $dir
      pid-check
      kill -$sig $pids
    fi
    # fuser will find open sockets that lsof won't, for example from gpg-agent.
    # note: -v shows kernel processes, which then doesn't return true when we want
    if pids=$(timeout 4 fuser -m $dir 2>/dev/null); then
      pid-check
      found_pids=true
      fuser -$sig -mvk $dir
    fi
    sleep .5
    if ! $found_pids; then
      return 0
    fi
  done
  return 1
}
umount-kill() {
  dir=$1
  if mountpoint -q $dir; then
    if m umount -R $dir; then
      unmounted+=($dir)
    else
      if ! kill-dir TERM TERM TERM INT INT HUP HUP TERM TERM TERM INT INT HUP HUP; then
        if $force; then  kill-dir KILL; fi
      fi

      if m umount -R $dir; then
        unmounted+=($dir)
      else
        echo "$0: failed to umount $dir"
        umount_ret=false
        ret=1
      fi
    fi
  fi
}

# duplicated in check-subvol
# Reassign $1 var from /dev/dm- to corresponding /dev/mapper/
mapper-dev() {
  local mapdev
  local -n devref=$1
  if [[ $devref == /dev/dm-* ]]; then
    for mapdev in /dev/mapper/*; do
      if [[ $(readlink -f $mapdev) == "$devref" ]]; then
        devref=$mapdev
        break
      fi
    done
  fi
}


##### begin command line parsing ########

# you can remove this if you do not have options which can have args with spaces or empty.

verbose=true
force=false
temp=$(getopt -l help,force,verbose hfv "$@") || usage 1
eval set -- "$temp"
while true; do
  case $1 in
    -f|--force) force=true ;;
    -v|--verbose) verbose=true ;;
    -h|--help) usage ;;
    --) shift; break ;;
    *) echo "$0: unexpected args: $*" >&2 ; usage 1 ;;
  esac
  shift
done

if (( $# )); then
  all_vols=( "$@" )
else
  all_vols=(q a o i ar qd qr)
fi

##### end command line parsing ########

ret=0

##### begin setup fstab for subvols we care about ######

if [[ -e /mnt/root/root2-crypttab ]]; then
  tu /etc/crypttab </mnt/root/root2-crypttab
  while read -r mapper_dev _; do
    if [[ ! -e /dev/mapper/$mapper_dev ]]; then
      m cryptdisks_start $mapper_dev
    fi
  done < <(cat /mnt/root/root2-crypttab)
fi
if [[ -e /mnt/root/root2-fstab ]]; then
  tu /etc/fstab </mnt/root/root2-fstab
  mnt /mnt/root2
  mnt /mnt/boot2
fi

root_dev=$(awk '$2 == "/" {print $1}' /etc/mtab)
mapper-dev root_dev
o_dev=$(awk '$2 == "/mnt/o" {print $1}' /etc/mtab)
mapper-dev o_dev


# root2_dev=$(awk '$2 == "/mnt/root2" {print $1}' /etc/mtab)
# mapper-dev root2_dev
# # dont bother with the above for crypt2_dev
# crypt2_dev=$root2_dev


if cryptsetup status $root_dev &>/dev/null; then
  crypt_dev=$root_dev
else # if we are in a recovery boot, find the next best crypt device
  mopts=,noauto
  # todo: I think I had an idea to not setup /o in this case,
  # but never finished implementing it
  for dev in $(dmsetup ls --target crypt | awk '{print $1}'); do
    dev=/dev/mapper/$dev
    if awk '{print $1}' /etc/mtab | grep -Fx $dev &>/dev/null; then
      crypt_dev=$dev
      break
    fi
  done
fi



# dont tax the cpus of old laptops
if (( $(nproc) > 2)); then
  mopts+=,compress=zstd
fi

fstab <<EOF
$crypt_dev  /a  btrfs  noatime,subvol=a$mopts  0 0
EOF

shopt -s nullglob

# ssh and probably some other things care about parent directory
# ownership, and ssh doesn\'t allow any group writable parent
# directories, so we are forced to use a directory structure similar
# to home directories
fa=(/mnt/root/btrbk/q.*); f=${fa[0]}
if [[ -e $f ]]; then
  fstab <<EOF
$crypt_dev  /q  btrfs  noatime,subvol=q$mopts  0 0
$crypt_dev  /qd  btrfs  noatime,subvol=qd$mopts  0 0
/q/p  /p  none  bind$mopts  0 0
EOF
fi

fa=(/mnt/root/btrbk/qr.*); f=${fa[0]}
if [[ -e $f ]]; then
  fstab <<EOF
$crypt_dev  /qr  btrfs  noatime,subvol=qr$mopts  0 0
EOF
fi

# not syncing ar at the moment
# fa=(/mnt/root/btrbk/ar.*); f=${fa[0]}
# if [[ -e $f ]]; then
#   fstab <<EOF
# $crypt_dev  /ar  btrfs  noatime,subvol=ar,uid=1000$mopts  0 0
# EOF
# fi


fa=(/mnt/o/btrbk/o.*); f=${fa[0]}
if [[ -e $f ]]; then
  if [[ $o_dev != "$root_dev" ]]; then
    # ,compress=zstd regardless of mopts since these are all text files
    # and it cuts disk use by about half.
    fstab <<EOF
$o_dev  /o  btrfs  noatime,subvol=o${mopts/,compress=zstd/},compress=zstd  0 0
EOF
  fi
  fstab <<EOF
/o/m  /m  none  bind$mopts  0 0
/o/debbugs  /debbugs  none  bind$mopts  0 0
EOF
fi


##### end setup fstab for subvols we care about ######

### begin get pids that this program depends on so we dont kill them
my_pids=($$ $PPID)
loop_limit=30
count=0
while [[ ${my_pids[-1]} != 1 && ${my_pids[-1]} != "${my_pids[-2]}" && $count -lt $loop_limit  ]]; do
  count=$((count + 1))
  p=$(ps -p ${my_pids[-1]} -o ppid=)
  if [[ $p == 0 || ! $p ]]; then
    break
  fi
  my_pids+=($p)
done
### end get pids that this program depends on so we dont kill them

for vol in ${all_vols[@]}; do
  d=/$vol
  if ! awk '$3 == "btrfs" {print $2}' /etc/fstab | grep -xF $d &>/dev/null; then
    continue
  fi


  ##### begin building up list of bind mounts  ######
  binds=() # list of bind mounts
  roots=($d)
  while true; do
    new_roots=()
    for r in ${roots[@]}; do
      # eg. when r=/q/p, for lines like
      # /q/p  /p  none  bind  0 0
      # output /p
      new_roots+=("$(sed -rn "s#^$r/\S+\s+(\S+)\s+none\s+(\S+,|)bind[[:space:],].*#\1#p" /etc/fstab)")
    done
    (( ${#new_roots} )) || break
    binds+=(${new_roots[@]})
    # roots is used to recursively find binds of binds if they exist.
    roots=( ${new_roots[@]} )
  done
  ##### end building up list of bind mounts  ######



  # if latest is already mounted, make sure binds are mounted and move on
  m check-subvol-stale $d
  # populated by check-subvol-stale if stale
  if ! fresh_snap=$(cat /nocow/btrfs-stale/$vol 2>/dev/null); then
    mnt $d
    did=$(stat -c%d $d)
    for b in ${binds[@]}; do
      if mountpoint -q $b; then
        bid=$(stat -c%d $b)
        if [[ $did != "$bid" ]]; then
          umount-kill $b
        fi
      fi
      mnt $b
    done
    continue
  fi


  ##### begin checking for loopback mounts ####
  found_loop=false
  for l in $(losetup -ln|awk '{print $6}'); do
    for dir in $d ${binds[@]}; do
      if [[ $l == $dir* ]]; then
        echo "$0: found loopback mount $l. giving up on unmounting $dir"
        ret=1
        found_loop=true
        break
      fi
    done
    if $found_loop; then
      break
    fi
  done
  if $found_loop; then
    continue
  fi
  ##### end end checking loopback mounts ####


  ## not using arbtt at the moment
  # if [[ $vol == q ]]; then
  #   ## allow to fail, user might not be logged in
  #   x sudo -u $(id -nu 1000) XDG_RUNTIME_DIR=/run/user/1000 systemctl --user stop arbtt ||:
  # fi
  umount_ret=true
  unmounted=()
  for dir in $(echo $d ${binds[*]}\ |tac -s\ ); do
    umount-kill $dir
  done

  # if we unmounted some but not all, restore them and move on
  if ! $umount_ret; then
    for dir in ${unmounted[@]}; do
      mnt $dir
    done
    continue
  fi

  #### begin dealing with leaf vols ####

  ### begin getting root_dir
  ### this is duplicated in check-subvol-stale

  dev=$(sed -rn "s,^\s*([^#]\S*)\s+$d\s.*,\1,p" /etc/fstab /etc/mtab|head -n1)
  d dev=$dev
  # note, we need $dev because $d might not be mounted, and we do this loop
  # because the device in fstab for the rootfs can be different.
  for devx in $(btrfs fil show $dev| sed -rn 's#.*path (\S+)$#\1#p'); do
    if [[ $devx == dm-* ]]; then
      devx=/dev/$devx
      mapper-dev devx
    fi
    d devx=$devx
    root_dir=$(sed -rn "s,^\s*$devx\s+(\S+).*\bsubvolid=[05]\b.*,\1,p" /etc/mtab /etc/fstab|head -n1)
    if [[ $root_dir ]]; then
      d root_dir=$root_dir
      break
    fi
  done
  if [[ ! $root_dir ]]; then
    echo "$0: error could not find root subvol mount for $dev" >&2
    exit 1
  fi
  ### end getting root_dir

  cd $root_dir
  if [[ -e $vol  ]]; then
    if [[ $vol == qd ]]; then
      m btrfs sub del qd
    else
      leaf=$vol.leaf.$(date +%Y-%m-%dT%H:%M:%S%z)
      m mv $vol $leaf
      m btrfs property set -ts $leaf ro true

      ### begin check if leaf is different, delete it if not ###
      parentid=$(btrfs sub show $leaf | awk '$1 == "Parent" && $2 == "UUID:" {print $3}')
      bsubs=(btrbk/$vol.*)
      bsub= # base subvolume
      # go in reverse order as its more likely to be at the end
      for ((i=${#bsubs[@]}-1; i>=0; i--)); do
        if [[ $parentid == $(btrfs sub show ${bsubs[i]} | awk '$1 == "UUID:" {print $2}') ]]; then
          bsub=${bsubs[i]}
          break
        fi
      done
      if [[ $bsub ]]; then
        # in testing, same subvol is 136 bytes. allow some overhead. 32 happens sometimes under systemd.
        # $ errno 32
        # EPIPE 32 Broken pipe
        lines=$(btrfs send --no-data -p $bsub $leaf | btrfs receive --dump | head -n 100 | wc -l || [[ $? == 141 || ${PIPESTATUS[0]} == 32 ]])
        if [[ $lines == 0 ]]; then
          # example output of no differences:
          # snapshot        ./qrtest                        uuid=c41ff6b7-0527-f34d-95ac-190eecf54ff5 transid=2239 parent_uuid=64949e1b-4a3e-3945-9a8e-cd7b7c15d7d6 parent_transid=2239
          echo suspected identical: $bsub $leaf
          x btrfs sub del $leaf
        fi
      fi
      ### end check if leaf is different, delete it if not ###

      ## begin expire leaf vols ##
      leaf_vols=($vol.leaf.*)
      count=${#leaf_vols[@]}
      leaf_limit_time=$(( EPOCHSECONDS - 60*60*24*60 )) # 60 days
      leaf_new_limit_time=$(( EPOCHSECONDS - 60*60*24 * 5 )) # 5 days this
      # goes backwards from oldest. leaf_new_limit_time is a safety
      # measure to ensure we don't delete very recent leafs.
      for leaf in ${leaf_vols[@]}; do
        leaf_time=$(date -d ${leaf#"$vol".leaf.} +%s)
        if (( leaf_limit_time > leaf_time || ( leaf_new_limit_time > leaf_time && count > 30 ) )); then
          x btrfs sub del $leaf
        fi
        count=$((count-1))
      done
    fi
    ## end expire leaf vols ##
  fi
  #### end dealing with leaf vols ####

  # Note, we make a few assumptions in this script, like
  # $d was not a different subvol id than $vol, and
  # things otherwise didn't get mounted very strangely.
  m btrfs sub snapshot $fresh_snap $vol
  for dir in $d ${binds[@]}; do
    m mnt $dir
  done

  ## arbtt disabled for now
  # if [[ $vol == q ]]; then
  #   # maybe this will fail if X is not running
  #   x sudo -u $(id -nu 1000) XDG_RUNTIME_DIR=/run/user/1000 systemctl --user start arbtt ||:
  # fi

  stale_dir=/nocow/btrfs-stale
  rm -f $stale_dir/$d
done



for dir in /mnt/r7/amy/{root/root,boot/boot}_ubuntubionic /mnt/{root2/root,boot2/boot}_ubuntubionic; do
  vol=${dir##*/}
  root_dir=${dir%/*}
  if [[ ! -d $root_dir ]]; then
    # this only exists on host kd currently
    continue
  fi
  # if latest is already mounted, make sure binds are mounted and move on
  m check-subvol-stale -p $dir
  # populated by check-subvol-stale if stale
  if ! fresh_snap=$(cat /nocow/btrfs-stale/$vol 2>/dev/null); then
    continue
  fi
  if [[ -d $dir ]]; then
    if ! kill-dir TERM TERM TERM INT INT HUP HUP TERM TERM TERM INT INT HUP HUP; then
      if $force; then  kill-dir KILL; fi
    fi
    m btrfs sub del $dir
  fi
  m btrfs sub snapshot $fresh_snap $dir
  rm -f /nocow/btrfs-stale/$vol
done

if (( ret >= 1 )); then
  echo "$0: exit status $ret. see error above"
fi
exit $ret