#!/bin/bash
# info: add system quota
# options: NONE
#
# example: v-add-sys-quota
#
# This function enables filesystem quota on /home partition
# Supports XFS and ext4 (both external and native quotas)

#----------------------------------------------------------#
#                 Variable & Function                      #
#----------------------------------------------------------#

# Includes
source /etc/tuliocp/tulio.conf
# shellcheck source=/usr/local/tulio/func/main.sh
source "$TULIO"/func/main.sh
# load config file
source_conf "$TULIO/conf/tulio.conf"

#----------------------------------------------------------#
#                    Install packages                      #
#----------------------------------------------------------#

# Install quota package
if ! type -P quota &> /dev/null; then
	export DEBIAN_FRONTEND=noninteractive
	apt-get -y install quota > /dev/null 2>&1
	check_result $? "quota package installation failed" "$E_UPDATE"
fi

# Perform verification if read-only mode is enabled
check_tulio_demo_mode

#----------------------------------------------------------#
#                       Action                             #
#----------------------------------------------------------#

# Extracting file system and partition info
mnt=$(df -P /home | awk '{print $6}' | tail -n1)
dev=$(df -P /home | awk '{print $1}' | tail -n1)
lnr=$(awk -v mnt="$mnt" '$1 !~ /^#/ && $2 == mnt {print NR; exit}' /etc/fstab)

# Verify line number was found
if [ -z "$lnr" ]; then
	log_history "Error: No entry found in /etc/fstab for $mnt"
	exit "$E_NOTEXIST"
fi

# Read fstab line details and create variables fstype and opt
read -r _ _ fstype opt _ _ < <(sed -n "${lnr}p" /etc/fstab)

# Policy for quotaon.service handling:
#   leave        -> don't touch the unit
#   mask         -> systemctl mask --now quotaon.service (recommended for ext4 native quotas and XFS)
#   unmask_start -> systemctl unmask + start quotaon.service (recommended for ext4 external quotas)
quotaon_policy="leave"

case "$fstype" in
	ext4)
		log_history "ext4 filesystem detected on $mnt. Configuring quota."

		# Detect native quota support
		fs_features=$(tune2fs -l "$dev" 2> /dev/null | grep "^Filesystem features:")
		native_quota_enabled=0

		if echo "$fs_features" | grep -q "quota"; then
			native_quota_enabled=1
			log_history "Native quota detected on ext4 filesystem."
		fi

		# Detect project feature (Only works with native quotas on ext4)
		project_feature_available=0
		if echo "$fs_features" | grep -q "project"; then
			project_feature_available=1
			log_history "Project feature detected on ext4 filesystem."
		fi

		# Check for external quota configuration in current mount options
		external_quota_configured=0
		if echo "$opt" | grep -qE "usrjquota=|grpjquota=|prjjquota=|jqfmt="; then
			external_quota_configured=1
			log_history "External quota configuration detected in fstab."
		fi

		# Determine quota options based on native support
		if [[ $native_quota_enabled -eq 1 ]]; then
			log_history "Using native ext4 quota support."

			# If switching from external to native, clean up first
			if [ $external_quota_configured -eq 1 ]; then
				log_history "Removing external quota configuration before enabling native quotas."
				clean_opts=$(echo "$opt" | tr ',' '\n' \
					| grep -vE '^(usrquota|grpquota|prjquota|usrjquota=|grpjquota=|prjjquota=|jqfmt=|uquota|gquota|pquota)' \
					| tr '\n' ',' | sed 's/,$//')

				if [ -n "$clean_opts" ]; then
					sed -i "${lnr}s|$opt|$clean_opts|" /etc/fstab
				else
					# If no other options, use defaults
					sed -i "${lnr}s|$opt|defaults|" /etc/fstab
				fi

				systemctl daemon-reload
				if mount -o remount "$mnt"; then
					log_history "Removed external quota options from fstab."
				else
					log_history "WARNING: Failed to remount after removing external quotas."
				fi

				# Update $opt to reflect current state
				opt="$clean_opts"
				[ -z "$opt" ] && opt="defaults"
			fi

			# Build native options: usr and grp are standard
			new_opts="usrquota,grpquota"

			# Add prjquota ONLY if native quota AND project feature are present
			if [[ $project_feature_available -eq 1 ]]; then
				new_opts="$new_opts,prjquota"
				log_history "Enabling native project quota support."
			fi

			# quotaon.service fails noisily on ext4 native quotas -> mask it
			quotaon_policy="mask"
		else
			log_history "Using external quota files for ext4 (Project quota not supported here)."

			# Ensure kernel modules are available for external quotas
			if ! find "/lib/modules/$(uname -r)" -type f -name '*quota_v*.ko*' | grep -q .; then
				log_history "Installing required kernel modules for quota support..."
				reboot_req="Y"
				apt-get -qq install linux-image-extra-virtual -y
				check_result $? "kernel module installation failed" "$E_UPDATE"
			fi

			# External quotas: only user and group
			new_opts="usrquota,grpquota,usrjquota=aquota.user,grpjquota=aquota.group,jqfmt=vfsv0"

			# For external quotas, allow quotaon.service to run at boot
			quotaon_policy="unmask_start"
		fi

		# Remove all quota-related options from current options
		clean_opts=$(echo "$opt" | tr ',' '\n' \
			| grep -vE '^(usrquota|grpquota|prjquota|usrjquota=|grpjquota=|prjjquota=|jqfmt=|uquota|gquota|pquota)' \
			| tr '\n' ',' | sed 's/,$//')

		# Update fstab if needed
		if ! echo "$opt" | grep -qF "$new_opts"; then
			if [ -n "$clean_opts" ]; then
				new_full_opts="$clean_opts,$new_opts"
			else
				new_full_opts="$new_opts"
			fi

			sed -i "${lnr}s|$opt|$new_full_opts|" /etc/fstab
			systemctl daemon-reload

			if mount -o remount "$mnt"; then
				log_history "Remounted $mnt with quota options: $new_opts"
			else
				log_history "WARNING: Failed to remount $mnt"
			fi
		fi

		# Create external quota files if using external quotas
		if [[ $native_quota_enabled -eq 0 ]]; then
			if [ ! -e "$mnt/aquota.user" ] || [ ! -e "$mnt/aquota.group" ]; then
				quotacheck -avcugm > /dev/null 2>&1
				log_history "Created external quota files."
			fi
		fi
		;;
	xfs)
		log_history "XFS filesystem detected on $mnt."

		# XFS quota options (always the same for XFS)
		xfs_quota_opts="uquota,gquota,pquota"

		# Check if quota options are already present
		if echo "$opt" | grep -qE "uquota|gquota|pquota"; then
			log_history "XFS quota options already present in mount options."
		else
			needs_update=1
		fi

		# XFS quotas are controlled by mount/rootflags; quotaon.service is unnecessary
		quotaon_policy="mask"

		# Handle root partition (requires GRUB modification)
		if [[ "$mnt" == "/" ]]; then
			log_history "/home is on root partition. Configuring XFS quota via GRUB and fstab."

			# Check and modify GRUB configuration
			grub_conf="/etc/default/grub"
			grub_conf="/etc/default/grub"
			if ! grep -qE "rootflags=.*(uquota|gquota|pquota)" "$grub_conf"; then
				# Add or update rootflags with uquota, gquota, pquota
				sed -i -E '
        /^GRUB_CMDLINE_LINUX=/{
          # If rootflags= already exists...
          /rootflags=/{
            /rootflags=[^" ]*uquota/! s/rootflags=([^" ]*)/rootflags=\1,uquota/
            /rootflags=[^" ]*gquota/! s/rootflags=([^" ]*)/rootflags=\1,gquota/
            /rootflags=[^" ]*pquota/! s/rootflags=([^" ]*)/rootflags=\1,pquota/
          }
          # If rootflags= does NOT exist, add it complete without touching the rest
          /rootflags=/! s/^(GRUB_CMDLINE_LINUX="[^"]*)"/\1 rootflags=uquota,gquota,pquota"/
        }
        ' "$grub_conf"
				check_result $? "Failed to add kernel parameters to $grub_conf"

				# Clean up duplicate spaces and spaces before closing quotes
				sed -i -E '/^GRUB_CMDLINE_LINUX=/ {s/  +/ /g; s/" /"/; s/ "/"/;}' "$grub_conf"
				check_result $? "Failed to clean up GRUB configuration"

				update-grub > /dev/null 2>&1
				log_history "GRUB updated with XFS quota flags. Reboot REQUIRED."
				reboot_req="Y"
			else
				log_history "XFS quota flags already present in GRUB configuration."
			fi

			# Update fstab for consistency (even though GRUB controls it)
			if [[ ${needs_update:-0} -eq 1 ]]; then
				# Remove any existing quota options first
				clean_opts=$(echo "$opt" | tr ',' '\n' \
					| grep -vE '^(uquota|gquota|pquota|usrquota|grpquota|prjquota)' \
					| tr '\n' ',' | sed 's/,$//')

				if [ -n "$clean_opts" ]; then
					new_full_opts="$clean_opts,$xfs_quota_opts"
				else
					new_full_opts="$xfs_quota_opts"
				fi

				sed -i "${lnr}s|$opt|$new_full_opts|" /etc/fstab
				check_result $? "Failed to modify fstab for XFS quota"
				systemctl daemon-reload
				log_history "XFS quota options added to fstab for root partition."
			fi

		# Handle separate partition (fstab only)
		else
			log_history "/home is on separate partition ($mnt). Configuring XFS quota via fstab."

			if [[ ${needs_update:-0} -eq 1 ]]; then
				# Remove any existing quota options first
				clean_opts=$(echo "$opt" | tr ',' '\n' \
					| grep -vE '^(uquota|gquota|pquota|usrquota|grpquota|prjquota)' \
					| tr '\n' ',' | sed 's/,$//')

				if [ -n "$clean_opts" ]; then
					new_full_opts="$clean_opts,$xfs_quota_opts"
				else
					new_full_opts="$xfs_quota_opts"
				fi

				sed -i "${lnr}s|$opt|$new_full_opts|" /etc/fstab
				check_result $? "Failed to modify fstab for XFS quota"
				systemctl daemon-reload

				if mount -o remount "$mnt"; then
					if quotaon -p "$dev" 2> /dev/null | grep -q '^user.*on$'; then
						log_history "Remounted $mnt with XFS quota options."
					else
						log_history "Reboot required to apply quota options on $mnt."
						reboot_req="Y"
					fi
				else
					log_history "WARNING: Failed to remount $mnt. Reboot required."
					reboot_req="Y"
				fi
			fi
		fi
		;;

	*)
		# Unsupported filesystem
		echo "Error: Tulio does not support quotas on $fstype filesystem."
		log_history "Quota setup failed: $fstype filesystem is not supported."
		exit "$E_DISK"
		;;
esac

#----------------------------------------------------------#
#                 quotaon.service handling                 #
#----------------------------------------------------------#
if systemctl list-unit-files quotaon.service > /dev/null 2>&1; then
	case "$quotaon_policy" in
		mask)
			log_history "Masking quotaon.service."
			systemctl mask --now quotaon.service > /dev/null 2>&1 || true
			;;
		unmask_start)
			log_history "Unmasking and starting quotaon.service."
			systemctl unmask quotaon.service > /dev/null 2>&1 || true
			systemctl start quotaon.service > /dev/null 2>&1 || true
			;;
		leave | *)
			log_history "Leaving quotaon.service unchanged."
			;;
	esac
fi

# Adding quotacheck on reboot (only for ext4 with external quotas)
if [[ "$fstype" == "ext4" ]] && [[ $native_quota_enabled -eq 0 ]]; then
	touch /forcequotacheck

	# Adding cron job for periodic quota check
	echo '#!/bin/bash' > /etc/cron.daily/quotacheck
	echo 'touch /forcequotacheck' >> /etc/cron.daily/quotacheck
	chmod a+x /etc/cron.daily/quotacheck
	log_history "Quotacheck cron job configured for external quotas."
else
	# Remove quotacheck cron job if it exists (not needed for XFS or ext4 native quotas)
	if [ -f /etc/cron.daily/quotacheck ]; then
		rm -f /etc/cron.daily/quotacheck
		log_history "Removed quotacheck cron job (not needed for $fstype)."
	fi

	# Remove aquota.* files if they exist (not needed for XFS or ext4 native quotas)
	quota_files=("quota.user" "quota.group" "quota.project" "aquota.user" "aquota.group" "aquota.project")
	for idx in "${quota_files[@]}"; do
		if [ -e "$mnt/$idx" ]; then
			if lsattr -d "$mnt/$idx" 2> /dev/null | awk '{print $1}' | grep -q 'i'; then
				log_history "INFO" "Removing immutable attribute from $mnt/$idx"
				chattr -i "$mnt/$idx"
			fi
			log_history "INFO" "Removing $mnt/$idx"
			rm -f "$mnt/$idx"
		fi
	done
fi

# Enabling group and user quota
if quotaon -pa | grep " $mnt " | grep -E 'user|group' | grep -q 'is off'; then
	quotaon -v "$mnt" 2> >(grep -vE 'Cannot stat.*device tmpfs|kernel probably supports ext4 quota' >&2) > /dev/null
	check_result $? "quota can't be enabled in $mnt" "$E_DISK"
fi

# Updating tulio.conf value
"$BIN"/v-change-sys-config-value "DISK_QUOTA" "yes"

# Rebuilding user quota (only if reboot is not required)
if [ "$reboot_req" != "Y" ]; then
	for user in $("$BIN"/v-list-users list); do
		"$BIN"/v-update-user-quota "$user"
	done
	log_history "User quotas updated successfully."
else
	# Create a one-time script to update quotas after reboot
	cat > /usr/local/tulio/bin/v-update-quotas-after-reboot.sh << 'EOF'
#!/bin/bash
BIN=/usr/local/tulio/bin
for user in $("$BIN/v-list-users" list); do
        "$BIN/v-update-user-quota" "$user"
done
# Self-delete after execution
rm -f /usr/local/tulio/bin/v-update-quotas-after-reboot.sh
rm -f /etc/cron.d/tulio-quota-reboot
EOF
	chmod +x /usr/local/tulio/bin/v-update-quotas-after-reboot.sh

	# Create one-time cron job
	cat > /etc/cron.d/tulio-quota-reboot << EOF
@reboot root sleep 30 && /usr/local/tulio/bin/v-update-quotas-after-reboot.sh
EOF
	chmod 644 /etc/cron.d/tulio-quota-reboot

	log_history "Created one-time job to update user quotas after reboot."
	echo "[ ! ] Note: User quotas will be automatically updated after the system reboots."
fi

#----------------------------------------------------------#
#                       Tulio                             #
#----------------------------------------------------------#

# Logging
"$BIN"/v-log-action "system" "Info" "Plugins" "System Quota enforcement enabled."
log_history "system quota enforcement enabled"
log_event "$OK" "$ARGUMENTS"

if [ "$reboot_req" = "Y" ]; then
	log_history "A system reboot is required to complete enable quota."
	echo "[ ! ] Note: A system reboot is required to complete quota setup."
fi

exit
