#!/usr/bin/bash

DAEMON_DIR="$DAEMON_DIR"

DIAG_DIR_BASE="/var/lib/vstorage/cs"
DIAG_DIR="$DIAG_DIR_BASE/$CLUSTER/$CS_ID"

RECOVERED_DIR="$DIAG_DIR_BASE/$CLUSTER/$CS_ID/recovered"
RECOVERED_LOGS="$RECOVERED_DIR/logs"
RECOVERED_TRASH="$RECOVERED_DIR/trash"
RECOVERED_JOURNAL_DIR="$RECOVERED_DIR/journal"
RECOVERED_JOURNAL_PART="$RECOVERED_JOURNAL_DIR/journal.part"
RECOVERED_JOURNAL="$RECOVERED_JOURNAL_DIR/journal"

UUID_FILE="$DIAG_DIR/uuid"
MOUNT_OPTS_FILE="$DIAG_DIR/mount_info"

CS_MODE_PATH="$DIAG_DIR/cs_mode"

CS_MODE_RW=0
CS_MODE_RO=1
CS_MODE_AUTO_REC=2
CS_MODE_MANUAL_REC=4
CS_MODE_FAIL_NO_REC=8
CS_MODE_SSD=16
CS_MODE_CSUM_DISABLED=32

FSCK_TIMEOUT=120

check_env()
{
	if [ -z "$CLUSTER" ]; then
		echo "Cluster name not specified" >&2
		return 1
	fi

	return 0
}

load_mode_flags()
{
	local mode=0

	ro_state=0
	auto_rec_state=0
	manual_rec_state=0
	on_ssd=0
	csum_disabled=0

	if [ ! -f "$CS_MODE_PATH" ]; then
		return 0
	fi

	if ! read -r mode < "$CS_MODE_PATH"; then
		echo "Failed to read CS mode" >&2
		return 1
	fi

	if ! [[ "$mode" =~ ^[0-9]+$ ]]; then
		echo "Invalid mode: $mode" >&2
		return 1
	fi

	ro_state=$(( (mode & CS_MODE_RO) != 0 ))
	auto_rec_state=$(( (mode & CS_MODE_AUTO_REC) != 0 ))
	manual_rec_state=$(( (mode & CS_MODE_MANUAL_REC) != 0 ))
	on_ssd=$(( (mode & CS_MODE_SSD) != 0 ))
	csum_disabled=$(( (mode & CS_MODE_CSUM_DISABLED) != 0 ))

	return 0
}

is_fs_ro()
{
	findmnt -no OPTIONS -T "$DAEMON_DIR" | tr ',' '\n' | grep -Fxq ro
}

clear_mode_flag_all()
{
	if ! rm -f -- "$CS_MODE_PATH"; then
		echo "failed to remove $CS_MODE_PATH" >&2
		return 1
	fi

	return 0
}

set_mode_flag()
{
	local flag="$1"
	local mode=0

	if [ -f "$CS_MODE_PATH" ]; then
		if ! read -r mode < "$CS_MODE_PATH"; then
			echo "Failed to read cs mode" >&2
			return 1
		fi
	fi

	if ! [[ "$mode" =~ ^[0-9]+$ ]]; then
		echo "Invalid mode in file: $mode" >&2

		if ! clear_mode_flag_all; then
			return 1
		fi
	fi

	mode=$(( mode | flag ))

	if ! echo "$mode" > "$CS_MODE_PATH"; then
		echo "Failed to write cs mode" >&2
		return 1
	fi

	return 0
}

clear_mode_flag()
{
	local flag="$1"
	local mode=0

	if [ -f "$CS_MODE_PATH" ]; then
		if ! read -r mode < "$CS_MODE_PATH"; then
			echo "Failed to read cs mode" >&2
			return 1
		fi
	fi

	if ! [[ "$mode" =~ ^[0-9]+$ ]]; then
		echo "Invalid mode in file: $mode" >&2
		return 1
	fi

	mode=$(( mode & ~flag ))

	if ! echo "$mode" > "$CS_MODE_PATH"; then
		echo "Failed to write cs mode" >&2
		return 1
	fi
}

read_uuid()
{
	if [ ! -r "$UUID_FILE" ]; then
		echo "UUID file is not readable" >&2
		return 1
	fi

	uuid=$(tr -d '\n' < "$UUID_FILE")
	if [ -z "$uuid" ]; then
		echo "UUID file is empty" >&2
		return 1
	fi
}

load_mount_info()
{
	mount_point=$(findmnt -no TARGET -T "$DAEMON_DIR" | tr -d '\n')
	if [ -z "$mount_point" ]; then
		echo "Failed to read mount point" >&2
		return 1
	fi

	mount_src=$(findmnt -no SOURCE -T "$DAEMON_DIR" | tr -d '\n')
	if [ -z "$mount_src" ]; then
		echo "Failed to read mount src" >&2
		return 1
	fi

	mount_opts=$(findmnt -no OPTIONS -T "$DAEMON_DIR" | tr -d '\n')
	if [ -z "$mount_opts" ]; then
		echo "Failed to read mount options" >&2
		return 1
	fi

	return 0
}

load_mount_info_from_dump()
{
	local to_discard=""

	if [ ! -f "$MOUNT_OPTS_FILE" ]; then
		echo "mount opts file not found" >&2
		return 1
	fi

	if ! read -r mount_point to_discard < "$MOUNT_OPTS_FILE"; then
		return 1
	fi

	if [ -z "$mount_point" ]; then
		echo "Failed to read mount point from dump" >&2
		return 1
	fi
}

terminate_process()
{
	local pid="$1"
	local timeout=30

	kill -15 $pid >& /dev/null

	while [ "$timeout" -ne 0 ]; do
		kill -0 $pid >& /dev/null
		if  [ $? -ne 0 ]; then
			return 0
		fi
		timeout=$((timeout-1))
		sleep 1
	done

	if kill -0 $pid 2>/dev/null; then
		timeout=30
		kill -9 $pid >& /dev/null

		while [ "$timeout" -ne 0 ]; do
			kill -0 $pid >& /dev/null
			if  [ $? -ne 0 ]; then
				return 0
			fi
			timeout=$((timeout-1))
			sleep 1
		done
	fi

	return 1
}

run_fsck()
{
	local timeout="$FSCK_TIMEOUT"

	e2fsck -fp "UUID=$uuid" &
	pid="$!"

	while [ "$timeout" -ne 0 ]; do
		kill -0 $pid >& /dev/null
		if [ $? -ne 0 ]; then
			break
		fi
		timeout=$((timeout-1))
		sleep 1
	done

	if kill -0 $pid 2>/dev/null; then
		# Still running after timeout
		if ! terminate_process $pid; then
			echo "Failed to terminate fsck" >&2
		fi
		return 1
	else
		wait "$pid"
		rc=$?

		if [ $rc -le 2 ]; then
			return 0
		else
			echo "fsck failed with '$rc'" >&2
			return 1
		fi
	fi
}

create_chunk_links()
{
	local i hex target_path link_path

	for i in {0..255}; do
		printf -v hex '%02x' "$i"

		target_path="$DAEMON_DIR/$hex"
		link_path="$RECOVERED_DIR/$hex"

		if ! ln -s -- "$target_path" "$link_path"; then
			echo "failed to create symlink: $link_path -> $target_path" >&2
			return 1
		fi
	done

	return 0
}

journal_fiemap_rec()
{
	# Using timeout 420. Overall this script should finish in 10 minutes.
	# In worst case scenario we will be 2 minutes in fsck,
	# then copying control dir, checksums and so on (20 sec max)
	# So we left with 7.5 minutes for journal recovery.
	# But if read is being proccessed at device level it can take
	# other 30-40 secs to stop. So using 7 minutes here.
	# Overall: 600 - (120 + 20 + 420 + 40) = 420
	if /usr/bin/vstorage -c "$CLUSTER" recover-cs -t 420 -e "$CS_ID"; then
		ln -s -- "$DIAG_DIR"/extracted.journal "$RECOVERED_JOURNAL" || return 1
		chown vstorage:vstorage "$DIAG_DIR"/extracted.journal || return 1
	else
		return 1
	fi
}

prepare_cs_env()
{
	local rc=0
	local real=$(realpath -m "$RECOVERED_DIR") || return 1

	case "$real" in
		"$DIAG_DIR_BASE"/*) ;;
		*)
			echo "Path not under $DIAG_DIR_BASE: $real" >&2
			return 1
			;;
	esac

	rm -rf -- "$real" || return 1

	install -d -m 0755 -o vstorage -g vstorage "$real" || return 1

	install -d -m 0755 -o vstorage -g vstorage "$RECOVERED_LOGS" || return 1
	install -d -m 0755 -o vstorage -g vstorage "$RECOVERED_TRASH" || return 1
	install -d -m 0755 -o vstorage -g vstorage "$RECOVERED_JOURNAL_DIR" || return 1

	rsync -a --delete -- "$DAEMON_DIR"/control "$RECOVERED_DIR/" || return 1

	rm -f -- "$RECOVERED_DIR"/control/journal || return 1
	ln -s -- "$RECOVERED_JOURNAL_DIR" "$RECOVERED_DIR"/control/journal || return 1

	ln -s -- "$DAEMON_DIR"/control/journal/checksum "$RECOVERED_DIR"/control/journal/checksum || return 1

	rsync -a --delete "$DAEMON_DIR"/control/journal/journal "$RECOVERED_JOURNAL_PART"
	rc=$?
	if [ "$rc" -eq 0 ]; then
		if ! mv -T -- "$RECOVERED_JOURNAL_PART" "$RECOVERED_JOURNAL"; then
			rm -f "$RECOVERED_JOURNAL_PART"
			journal_fiemap_rec || return 1
		fi

	else
		rm -f "$RECOVERED_JOURNAL_PART"
		journal_fiemap_rec || return 1
	fi

	create_chunk_links || return 1

	return 0
}

cs_put_event()
{
	local event_msg="$1"
	local rc=0

	/usr/bin/vstorage -q -c "$CLUSTER" put-event -e -s cs "$event_msg"
	rc=$?
	if [ "$rc" -ne 0 ]; then
		echo "Failed to write event log, rc=$rc" >&2
	fi

	echo "$event_msg"
}

cs_to_man_rec_mode()
{
	local event_msg="$1"
	local rc=0

	cs_put_event "CS#$CS_ID needs manual recovery. $event_msg"

	/usr/bin/vstorage -q -c "$CLUSTER" rm-cs -F "$CS_ID"
	rc=$?
	if [ "$rc" -ne 0 ]; then
		/usr/bin/vstorage -q -c "$CLUSTER" -f rm-cs -F "$CS_ID"
		rc=$?
		if [ "$rc" -ne 0 ]; then
			echo "Failed to put CS into failed mode on MDS, rc=$rc" >&2
		fi
	fi

	set_mode_flag "$CS_MODE_MANUAL_REC"
}

# For ext4 backup superblocks are located
# at the first block of the block group (except for group 0)
# With sparse_super (the default for many years), backups are only in
# group 1 and in groups whose group number is a power of 3, 5, or 7.
# Here first 8 are hardcoded, assuming
# 4 KiB block size, 32768 blocks per group.
try_mount_with_reserved_sb()
{
	local mounted=0

	while read -r sb; do
		sb_for_mount=$((sb * 4))   # for 4K ext4 block size
		echo "Trying backup superblock $sb_for_mount..." >&2

		if mount -o "ro,noload,noatime,sb=$sb_for_mount" "$mount_src" "$mount_point"; then
			echo "Mounted using backup superblock $sb_for_mount"
			mounted=1
			break
		fi
	done <<'EOF'
32768
98304
163840
229376
294912
819200
884736
1605632
EOF

	return $(( ! mounted ))
}

mount_full_ro()
{
	local mounted=0

	if ! blockdev --setro "$mount_src"; then
		cs_put_event "CS#$CS_ID Failed to switch repo partition $mount_src to ro mode."
		return 1
	fi

	if mount -o "ro,noload,noatime" "$mount_src" "$mount_point"; then
		mounted=1
	elif try_mount_with_reserved_sb; then
		mounted=1
	fi

	if [ $mounted -eq 1 ]; then
		if ! prepare_cs_env; then
			cs_to_man_rec_mode "Failed to prepare recovery environment."
			if ! umount -- "$mount_point"; then
				echo "Failed to unmount $mount_point" >&2
			fi

			return 1
		else
			set_mode_flag "$CS_MODE_AUTO_REC"
			return 0
		fi
	else
		cs_to_man_rec_mode "Failed to mount repo partition $mount_src in ro mode."
		return 1
	fi
}

disk_present_by_uuid()
{
	local uuid=$1

	if [ -z "$uuid" ]; then
		echo "Empty repo disk UUID"
		return 1
	fi

	lsblk -no UUID | grep -Fxq -- "$uuid"
}

main()
{
	local uuid=""
	local ro_state=0
	local auto_rec_state=0
	local manual_rec_state=0
	local on_ssd=0
	local csum_disabled=0
	local mount_point=""
	local mount_src=""
	local mount_opts=""
	local mount_fs_type=""

	# Backward compatibility
	if [ -z "$CS_ID" ]; then
		echo "CS_ID name not specified" >&2
		return 0
	fi

	if ! check_env; then
		return 1
	fi

	if [ ! -d "$DIAG_DIR" ]; then
		if ! /usr/bin/vstorage -c "$CLUSTER" recover-cs -b "$DAEMON_DIR"; then
			echo "Failed to create journal blocks for CS $SERVICEID ($DAEMON_DIR)" >&2
			return 0
		fi
	fi

	if ! load_mode_flags; then
		echo "Manual recovery required"
		return 1
	fi

	if is_fs_ro; then
		set_mode_flag "$CS_MODE_RO"
		ro_state=1
	fi

	if [ "$ro_state" -eq 0 ] && [ "$auto_rec_state" -eq 0 ] && [ "$manual_rec_state" -eq 0 ]; then
		return 0
	fi

	# Presence of DIAG dir and absence of UUID file
	# means cs_repo and journal were located
	# on different physical disks.
	# Recovery is not available.
	if [ ! -f "$UUID_FILE" ]; then
		echo "UUID file not present" >&2
		return 0
	fi

	if ! load_mount_info; then
		if ! load_mount_info_from_dump; then
			echo "Manual recovery required"
			return 1
		fi
		mount_opts="noatime,lazytime,nofail"
	fi

	# Docker workaround
	if [ "$mount_src" = "overlay" ]; then
		echo "Virtual environment. No recovery available" >&2

		if ! clear_mode_flag_all; then
			return 1
		fi

		return 0
	fi

	if [ "$manual_rec_state" -eq 1 ]; then
		cs_put_event "CS#$CS_ID Manual recovery required."
		return 1
	fi

	if [ "$on_ssd" -eq 1 ]; then
		cs_to_man_rec_mode "Recovery with journal on SSD is not supported"
		set_mode_flag "$CS_MODE_FAIL_NO_REC"
		return 1
	fi

	if [ "$csum_disabled" -eq 1 ]; then
		cs_to_man_rec_mode "CS#$CS_ID Recovery with journal without checksums enabled is not supported"
		set_mode_flag "$CS_MODE_FAIL_NO_REC"
		return 1
	fi

	if [ "$auto_rec_state" -eq 1 ]; then
		return 0
	fi

	if [ "$ro_state" -eq 0 ]; then
		return 0
	fi

	if ! read_uuid; then
		cs_to_man_rec_mode "Failed to read saved UUID."
		return 1
	fi

	if ! disk_present_by_uuid "$uuid"; then
		cs_put_event "CS#$CS_ID repository partition with UUID=$uuid was not found"
		return 1
	fi


	if findmnt -rno TARGET "$mount_point" >/dev/null 2>&1; then
		if ! umount -- "$mount_point"; then
			cs_put_event "CS#$CS_ID Unable to umount repo partition $mount_point."
			return 1
		fi
	fi

	# mount_src(dev_name) could've been changed on disk disconnect
	if ! mount_src=$(blkid -U "$uuid"); then
		cs_put_event "CS#$CS_ID Failed to resolve UUID $uuid to dev name" >&2
		return 1
	fi

	if [ -z "$mount_src" ]; then
		cs_put_event "CS#$CS_ID mount_src from $uuid resolved to empty string" >&2
		return 1
	fi

	if findmnt -rn -S "UUID=$uuid"; then
		echo "$mount_src is mounted"

		mnt=$(findmnt -rn -S "$mount_src" -o TARGET)
		echo "Mountpoint: $mnt"

		echo "Processes using the mount:"
		fuser -vm "$mnt"

		cs_put_event "CS#$CS_ID Partitition is mounted, unable to proceed $mount_point."
		return 1
	fi

	if ! blockdev --setrw "$mount_src"; then
		cs_put_event "CS#$CS_ID Failed to switch repo partition $mount_src to rw mode."
		return 1
	fi

	if run_fsck; then
		if mount -o "${mount_opts},rw" "$mount_src" "$mount_point"; then
			if ! /usr/bin/vstorage -c "$CLUSTER" -f rm-cs -U "$CS_ID"; then
				return 1
			fi

			if ! clear_mode_flag_all; then
				return 1
			fi

			return 0
		else
			if ! mount_full_ro; then
				return 1
			fi
		fi
	else
		if ! disk_present_by_uuid "$uuid"; then
			cs_put_event "CS#$CS_ID repository partition with UUID=$uuid was not found"
			return 1
		fi

		if ! mount_full_ro; then
			return 1
		fi
	fi

	return 0
}

main "$@"
