#!/bin/bash -eu

# this function is also called by archiveloop by sourcing this script
function read_setup_variables {
  if [ -z "${setup_file+x}" ]
  then
    local -r setup_file=/root/teslausb_setup_variables.conf
  fi
  if [ -e $setup_file ]
  then
    # "shellcheck" doesn't realize setup_file is effectively a constant
    # shellcheck disable=SC1090
    source $setup_file
  else
    echo "couldn't find $setup_file"
    return 1
  fi

  # TODO: change this "declare" to "local" when github updates
  # to a newer shellcheck.
  declare -A newnamefor

  newnamefor[archiveserver]=ARCHIVE_SERVER
  newnamefor[camsize]=CAM_SIZE
  newnamefor[musicsize]=MUSIC_SIZE
  newnamefor[sharename]=SHARE_NAME
  newnamefor[shareuser]=SHARE_USER
  newnamefor[sharepassword]=SHARE_PASSWORD
  newnamefor[tesla_email]=TESLA_EMAIL
  newnamefor[tesla_password]=TESLA_PASSWORD
  newnamefor[tesla_vin]=TESLA_VIN
  newnamefor[timezone]=TIME_ZONE
  newnamefor[usb_drive]=USB_DRIVE

  local oldname
  for oldname in "${!newnamefor[@]}"
  do
    local newname=${newnamefor[$oldname]}
    if [[ -z ${!newname+x} ]] && [[ -n ${!oldname+x} ]]
    then
      local value=${!oldname}
      export $newname="$value"
    fi
  done

  # set defaults for things not set in the config
  REPO=${REPO:-marcone}
  BRANCH=${BRANCH:-main-dev}
  CONFIGURE_ARCHIVING=${CONFIGURE_ARCHIVING:-true}
  UPGRADE_PACKAGES=${UPGRADE_PACKAGES:-false}
  export TESLAUSB_HOSTNAME=${TESLAUSB_HOSTNAME:-teslausb}
  SAMBA_ENABLED=${SAMBA_ENABLED:-false}
  SAMBA_GUEST=${SAMBA_GUEST:-false}
  SNAPSHOTS_ENABLED=${SNAPSHOTS_ENABLED:-true}
  INCREASE_ROOT_SIZE=${INCREASE_ROOT_SIZE:-0}
  export CAM_SIZE=${CAM_SIZE:-90%}
  export MUSIC_SIZE=${MUSIC_SIZE:-0}
  export USB_DRIVE=${USB_DRIVE:-''}
}

if [ "$0" == "-bash" ]
then
  # called from .bashrc as part of root login, set completion
  complete -W "diagnose upgrade selfupdate install" setup-teslausb
  return 0
fi

if [ "${BASH_SOURCE[0]}" != "$0" ]
then
  echo "${BASH_SOURCE[0]} was sourced, setting setup variables"
  read_setup_variables
  return 0 # shouldn't use exit when sourced
fi

if [ "${FLOCKED:-}" != "$0" ]
then
  if FLOCKED="$0" flock -en -E 99 "$0" "$0" "$@" || case "$?" in
  99) echo already running
      exit 99
      ;;
  *)  exit $?
      ;;
  esac
  then
    # success
    exit 0
  fi
fi

function setup_progress () {
  local setup_logfile=/boot/teslausb-headless-setup.log
  if [ -w $setup_logfile ]
  then
    echo "$( date ) : $*" >> "$setup_logfile"
  fi
  echo "$@"
}

if [[ $EUID -ne 0 ]]
then
  setup_progress "STOP: Run sudo -i."
  exit 1
fi

read_setup_variables

function dehumanize () {
  echo $(($(echo "$1" | sed 's/GB/G/;s/MB/M/;s/KB/K/;s/G/*1024M/;s/M/*1024K/;s/K/*1024/')))
}

REBOOT=false

# wpa_supplicant should have a country code. Use US as the default
# to avoid using the disallowed channels in the US.
WPA=/etc/wpa_supplicant/wpa_supplicant.conf
if ! grep -q "country=" $WPA
then
  setup_progress "adding country code to wpa_supplicant.conf"
  echo "country=US" >> $WPA
  REBOOT=true
fi

readonly BOOT_DEVICE=$(lsblk -dpno pkname /dev/disk/by-label/boot)
export BOOT_DEVICE
readonly BOOT_DEVICE_PART=$(lsblk -dpno name /dev/disk/by-label/boot | sed 's/1$//')
export BOOT_DEVICE_PART
INCREASE_ROOT_SIZE=$(($(dehumanize "$INCREASE_ROOT_SIZE") / 512))

if [ "$INCREASE_ROOT_SIZE" != "0" ] && [ ! -e "${BOOT_DEVICE_PART}3" ]
then
  if [ ! -e /root/TESLAUSB_ROOT_PARTITION_INCREASED ]
  then
    touch /root/TESLAUSB_ROOT_PARTITION_INCREASED
    ROOTSTART=$(partx --show -g -o START "${BOOT_DEVICE_PART}2")
    ROOTSIZE=$(partx --show -g -o SECTORS "${BOOT_DEVICE_PART}2")
    ROOTSIZE=$((ROOTSIZE + INCREASE_ROOT_SIZE))
    echo "$ROOTSTART,$ROOTSIZE" | sfdisk --force "${BOOT_DEVICE}" -N 2
    setup_progress "increased root partition size"
    REBOOT=true
  else
    setup_progress "increasing root filesystem size to match partition size"
    resize2fs "${BOOT_DEVICE_PART}2"
  fi
fi

if [ "$REBOOT" = "true" ]
then
  if [ -t 0 ]
  then
    setup_progress "please reboot for changes to take effect"
    exit
  else
    setup_progress "rebooting for changes to take effect"
    reboot
    exit
  fi
fi

function headless_setup_mark_setup_success () {
  rm -f /boot/TESLAUSB_SETUP_FAILED
  rm -f /boot/TESLAUSB_SETUP_STARTED
  touch /boot/TESLAUSB_SETUP_FINISHED
}

function isPi4 {
  grep -q "Pi 4" /sys/firmware/devicetree/base/model
}

function flash () {
  local ON=0
  local OFF=1
  if isPi4
  then
    ON=1
    OFF=0
  fi
  echo none > /sys/class/leds/led0/trigger
  echo $OFF > /sys/class/leds/led0/brightness
  sleep 1
  for ((i=1; i<=$1; i++))
  do
    echo $ON > /sys/class/leds/led0/brightness
    sleep .2
    echo $OFF > /sys/class/leds/led0/brightness
    sleep .8
  done
}

function headless_setup_progress_flash () {
  if [ ! -t 0 ]
  then
    flash "$1"
  fi
}

function verify_configuration () {
  get_script /tmp verify-configuration.sh setup/pi

  /tmp/verify-configuration.sh
}

function curlwrapper () {
  setup_progress "curl $*"
  while ! curl --fail "$@"
  do
    setup_progress "'curl $*' failed, retrying" > /dev/null
    sleep 3
  done
}

function get_script () {
  local local_path="$1"
  local name="$2"
  local remote_path="${3:-}"

  if [ -z ${CURL_PREFIX:+x} ]
  then
    HEAD_VERSION=$(curlwrapper -s "https://api.github.com/repos/$REPO/teslausb/git/refs/heads/$BRANCH" | grep '"sha":' | awk '{print $2}' | sed 's/"\|,//g')
    BRANCH="$HEAD_VERSION"
    export CURL_PREFIX=https://raw.githubusercontent.com/"$REPO"/teslausb/"$BRANCH"
  fi

  # shellcheck disable=SC2086
  curlwrapper -o "$local_path/$name" $CURL_PREFIX"/$remote_path/$name"
  chmod +x "$local_path/$name"
  setup_progress "Downloaded $local_path/$name ..."
}

function get_ancillary_setup_scripts () {
  get_script /tmp create-backingfiles-partition.sh setup/pi
  get_script /tmp create-backingfiles.sh setup/pi
  get_script /tmp make-root-fs-readonly.sh setup/pi
  get_script /tmp configure.sh setup/pi
}

function get_common_scripts () {
  get_script /root/bin remountfs_rw run
  get_script /root/bin make_snapshot.sh run
  get_script /root/bin mount_snapshot.sh run
  get_script /root/bin release_snapshot.sh run
  get_script /root/bin force_sync.sh run
}

function fix_cmdline_txt_modules_load ()
{
  setup_progress "Fixing the modules-load parameter in /boot/cmdline.txt..."
  cp /boot/cmdline.txt ~
  sed 's/ modules-load=dwc2,g_ether/ modules-load=dwc2/' ~/cmdline.txt > /boot/cmdline.txt
  rm ~/cmdline.txt
  setup_progress "Fixed cmdline.txt."
}

BACKINGFILES_MOUNTPOINT=/backingfiles
MUTABLE_MOUNTPOINT=/mutable

function create_usb_drive_backing_files () {
  if [ ! -e "$BACKINGFILES_MOUNTPOINT" ]
  then
    mkdir "$BACKINGFILES_MOUNTPOINT"
  fi

  if [ ! -e "$MUTABLE_MOUNTPOINT" ]
  then
    mkdir "$MUTABLE_MOUNTPOINT"
  fi

  /tmp/create-backingfiles-partition.sh "$BACKINGFILES_MOUNTPOINT" "$MUTABLE_MOUNTPOINT"

  if ! findmnt --mountpoint $BACKINGFILES_MOUNTPOINT
  then
    setup_progress "Mounting the partition for the backing files..."
    mount $BACKINGFILES_MOUNTPOINT
    setup_progress "Mounted the partition for the backing files."
  fi

  if [ ! -e $BACKINGFILES_MOUNTPOINT/cam_disk.bin ]
  then
    setup_progress "Creating backing disk files."
    /tmp/create-backingfiles.sh "$CAM_SIZE" "$MUSIC_SIZE" "$BACKINGFILES_MOUNTPOINT"
  else
    # mount cam image and make sure the right directories exist
    umount /mnt/cam || true
    if modprobe -r g_mass_storage && mount /mnt/cam
    then
      mkdir -p /mnt/cam/TeslaCam
      mkdir -p /mnt/cam/TeslaTrackMode
      umount /mnt/cam
    else
      setup_progress "STOP: Couldn't check image"
      exit 1
    fi
  fi
}

function configure_hostname () {
  local new_host_name="$TESLAUSB_HOSTNAME"
  local old_host_name
  old_host_name=$(hostname)

  # Set the specified hostname if it differs from the current name
  if [ "$new_host_name" != "$old_host_name" ]
  then
    setup_progress "Configuring the hostname..."
    sed -i -e "s/$old_host_name/$new_host_name/g" /etc/hosts
    sed -i -e "s/$old_host_name/$new_host_name/g" /etc/hostname
    while ! hostnamectl set-hostname "$new_host_name"
    do
      setup_progress "hostnamectl failed, retrying"
      sleep 1
    done
    systemctl restart avahi-daemon
    setup_progress "Configured hostname: $(hostname)"
  fi
}

function make_root_fs_readonly () {
  /tmp/make-root-fs-readonly.sh
}

function update_package_index () {
  setup_progress "Updating package index files..."
  while ! apt-get update
  do
    setup_progress "Failed, retrying"
    sleep 2
  done
}

function upgrade_packages () {
  if [ "$UPGRADE_PACKAGES" = true ]
  then
    setup_progress "Upgrading installed packages..."
    # clean the cache to free up space, since especially
    # a kernel update requires quite a bit of temporary
    # extra space
    apt-get clean
    apt-get --assume-yes upgrade
  else
    setup_progress "Skipping package upgrade."
  fi
  # no real need to keep the cache around after setup
  apt-get clean
  fstrim / || true
}

function set_timezone () {
  if [ -n "${TIME_ZONE:+x}" ]
  then
    if [ -f "/usr/share/zoneinfo/$TIME_ZONE" ]
    then
      ln -sf "/usr/share/zoneinfo/$TIME_ZONE" /etc/localtime
    elif [ "$TIME_ZONE" = "auto" ]
    then
      if curlwrapper -o /root/bin/tzupdate.py https://raw.githubusercontent.com/marcone/tzupdate/develop/tzupdate.py
      then
        apt-get -y --force-yes install python-requests
        chmod +x /root/bin/tzupdate.py
        if ! tzout=$(/root/bin/tzupdate.py 2>&1)
        then
          setup_progress "auto timezone failed: $tzout"
        else
          setup_progress "$tzout"
        fi
      fi
    else
      setup_progress "invalid timezone: $TIME_ZONE"
    fi
  fi
}

function cmd_install {
  /root/bin/remountfs_rw
  get_script /root/bin "$(basename "$1")" "$(dirname "$1")"
  setup_progress "$1 installed in /root/bin/"
  exit
}

function cmd_selfupdate {
  get_script /tmp setup-teslausb setup/pi &> /dev/null
  if cmp -s /tmp/setup-teslausb /root/bin/setup-teslausb
  then
    setup_progress "$0 already up to date"
  else
    /root/bin/remountfs_rw > /dev/null
    mv /tmp/setup-teslausb /root/bin/setup-teslausb
    setup_progress "$0 updated"
  fi
  setup_progress "other scripts may need to be updated by running $0"
  exit 0
}

function cmd_upgrade {
  if [ ! -e /boot/TESLAUSB_SETUP_FINISHED ]
  then
    echo "STOP: previous setup didn't finish, can't upgrade unfinished install"
    exit 1
  fi
  get_script /tmp setup-teslausb setup/pi &> /dev/null
  exec bash -c "/tmp/setup-teslausb upgrade_prepare && /tmp/setup-teslausb && /tmp/setup-teslausb upgrade_finish" < /dev/null
}

function cmd_upgrade_prepare {
  setup_progress "preparing for full upgrade"
  killall archiveloop || true
  service smbd stop || true
  service autofs stop || true
  umount /backingfiles/snapshots/snap*/mnt || true
  umount /mnt/cam /mnt/music || true
  umount /mnt/archive /mnt/musicarchive || true
  modprobe -r g_mass_storage
  /root/bin/remountfs_rw > /dev/null
}

function cmd_upgrade_finish {
  mv /tmp/setup-teslausb /root/bin/setup-teslausb
  setup_progress "upgrade finished"
  for i in {5..1}
  do
    echo -e -n "rebooting in $i seconds to apply changes, press ctrl-c to abort\r"
    sleep 1
  done
  echo -e '\nRebooting'
  systemctl start reboot.target
}

function cmd_diagnose {
  local hardware
  local os

  hardware=$( tr -d '\000' < /sys/firmware/devicetree/base/model )
  os=$(. /etc/os-release && echo "$PRETTY_NAME")
  {
    echo -e "====== summary ======"
    echo -e "hardware: ${hardware}"
    echo -e "OS: ${os}"
    if [ "${ARCHIVE_SYSTEM:-none}" = "cifs" ]
    then
      if grep -q '/mnt/archive' /etc/fstab
      then
        echo "CIFS archiving selected"
      else
        echo "ERROR: CIFS archiving selected, but archive not defined in fstab"
      fi
    else
      echo "archive method: ${ARCHIVE_SYSTEM:-none}"
    fi

    if ! blkid -L backingfiles > /dev/null
    then
      echo "ERROR: backingfiles partition does not exist"
    fi
    if [ ! -d /backingfiles ]
    then
      echo "backingfiles directory does not exist"
    fi
    if ! grep -q '/backingfiles' /etc/fstab
    then
      echo "ERROR: backingfiles not in fstab"
    fi

    if [ ! -f /backingfiles/cam_disk.bin ]
    then
      echo "ERROR: cam disk image does not exist"
    fi
    if ! grep -q '/backingfiles/cam_disk.bin' /etc/fstab
    then
      echo "ERROR: cam disk image not in fstab"
    fi
    for LUN0 in /sys/devices/platform/soc/??980000.usb/gadget/lun0/file
    do
      if [ -e "$LUN0" ]
      then
        echo "lun0 connected, from file $(cat "$LUN0")"
      fi
    done
    for LUN1 in /sys/devices/platform/soc/??980000.usb/gadget/lun1/file
    do
      if [ -e "$LUN1" ]
      then
        echo "lun1 connected, from file $(cat "$LUN1")"
      fi
    done
    if ! blkid -L mutable > /dev/null
    then
      echo "ERROR: mutable partition does not exist"
    fi
    if [ ! -d /mutable ]
    then
      echo "ERROR: mutable directory does not exist"
    fi
    if ! grep -q '/mutable' /etc/fstab
    then
      echo "ERROR: mutable not in fstab"
    fi

    numsnapshots=$( mount | grep -c snapshot )
    echo "$numsnapshots snapshots mounted"

    if [ ! -e /boot/TESLAUSB_SETUP_FINISHED ]
    then
      echo 'ERROR: setup did not finish'
    fi

    echo -e "====== disk / images ======"
    parted -s "${BOOT_DEVICE}" print || true
    if [ -n "${USB_DRIVE:+x}" ]
    then
      parted -s "$USB_DRIVE" print || true
    fi
    if [ -f /backingfiles/cam_disk.bin ]
    then
      echo "cam disk image has $(filefrag /backingfiles/cam_disk.bin | awk '{print $2}') extents"
      parted -s /backingfiles/cam_disk.bin print || true
    else
      echo "no cam disk image found"
    fi
    if [ -f /backingfiles/music_disk.bin ]
    then
      echo "music disk image has $(filefrag /backingfiles/music_disk.bin | awk '{print $2}') extents"
      parted -s /backingfiles/music_disk.bin print || true
    else
      echo "no music disk image found"
    fi

    echo -e "====== network ======"
    ifconfig
    iwconfig wlan0 | grep Link

    echo -e "====== fstab ======"
    if [ -e /etc/fstab ]
    then
      cat /etc/fstab
    else
      echo "no fstab found"
    fi

    echo -e "====== initial setup boot log ======"
    mkdir /tmp/root$$
    mount --bind / /tmp/root$$
    if [ -e /tmp/root$$/var/log/boot.log ]
    then
      cat /tmp/root$$/var/log/boot.log
    else
      echo "no boot log found"
    fi
    umount /tmp/root$$
    rmdir /tmp/root$$

    echo -e "====== setup log ======"
    if [ -e /boot/teslausb-headless-setup.log ]
    then
      cat /boot/teslausb-headless-setup.log
    else
      echo "no setup log found"
    fi

    echo -e "====== archiveloop log ======"
    if [ -e /mutable/archiveloop.log ]
    then
      cat /mutable/archiveloop.log
    else
      echo "no archiveloop log found"
    fi

    echo -e "====== system log ======"
    if [ -x /bin/logread ]
    then
      /bin/logread
    else
      echo "logread not installed"
    fi

    echo -e "====== dmesg ======"
    dmesg -T
    echo -e "====== process list and uptime ======"
    ps -eaf
    echo "$(hostname) has been $(uptime -p). System time is $(date)"
    echo -e "====== end of diagnostics ======"
  } |
    # clean up the output a bit
    tr '\r' '\n' |
    sed '/^ *$/d' |
    grep -a -v '^Reading package lists' |
    grep -a -v '^(Reading database' |
    grep -a -v "^Adding 'diversion of" |
    grep -a -v "^Removing 'diversion of" |
    sed -E 's/\o033\[0;32m//' |
    sed -E 's/\o033\[0m//'
}

export -f setup_progress

INSTALL_DIR=${INSTALL_DIR:-/root/bin}
if [ "$INSTALL_DIR" != "/root/bin" ]
then
  setup_progress "WARNING: 'INSTALL_DIR' setup variable no longer supported"
fi

BRANCHNAME="$BRANCH"

if [ -n "${1:+x}" ]
then
  command=cmd_$1
  if typeset -f "$command" > /dev/null
  then
    shift
    $command "$@"
    exit 0
  else
    setup_progress "unknown command: $1"
    exit 1
  fi
fi

# Update /boot/config.txt if needed
if ! grep -q 'dtoverlay=dwc2' /boot/config.txt
then
  echo -e "dtoverlay=dwc2\n" >> /boot/config.txt
fi

configure_hostname

tmpdir=/tmp/$$
mkdir -p "$tmpdir"
get_script "$tmpdir" setup-teslausb setup/pi &> /dev/null
if cmp -s "$tmpdir/setup-teslausb" "$0"
then
  setup_progress "$0 is up to date"
else
  setup_progress "WARNING: $BRANCHNAME contains a different version of $0. It is recommended to run '$0 selfupdate' to update to that version"
fi

update_package_index

# set time zone so we get decent timestamps in the rest of the setup log
set_timezone

# Flash for stage 2 headless (verify requested configuration)
headless_setup_progress_flash 2

setup_progress "Verifying that the requested configuration is valid..."

verify_configuration

# Flash for Stage 3 headless (grab scripts)
headless_setup_progress_flash 3

setup_progress "Downloading additional setup scripts."

mkdir -p /root/bin

get_common_scripts

get_ancillary_setup_scripts

pushd ~

fix_cmdline_txt_modules_load

# Flash for stage 4 headless (Create backing files)
headless_setup_progress_flash 4

create_usb_drive_backing_files

if [ "$CONFIGURE_ARCHIVING" = true ]
then
  setup_progress "calling configure.sh"
  export -f curlwrapper
  export -f get_script
  /tmp/configure.sh
else
  setup_progress "skipping configure.sh"
fi

if [ "$SAMBA_ENABLED" = "true" ]
then
  if [ "$SNAPSHOTS_ENABLED" = "true" ]
  then
    export SAMBA_GUEST
    get_script /tmp configure-samba.sh setup/pi
    /tmp/configure-samba.sh
  else
    setup_progress "snapshots disabled, skipping Samba configuration"
  fi
fi

if [ -n "${AP_SSID:+x}" ]
then
  get_script /tmp configure-ap.sh setup/pi
  /tmp/configure-ap.sh
fi

if [ "$SNAPSHOTS_ENABLED" = "true" ]
then
  get_script /tmp configure-automount.sh setup/pi
  /tmp/configure-automount.sh
else
  setup_progress "skipping automount setup because snapshots are disabled"
fi

# source setup-teslausb from .bashrc to set up completion
if ! grep -q setup-teslausb /root/.bashrc
then
  echo "source /root/bin/setup-teslausb" >> /root/.bashrc
fi

make_root_fs_readonly

upgrade_packages

headless_setup_mark_setup_success

# Flash for stage 5 headless (Mark success, FS readonly)
headless_setup_progress_flash 5

setup_progress "All done."

if [ -t 0 ]
then
  setup_progress '(reboot now for changes to take effect)'
fi