#!/usr/bin/env bash
# NAME: Configure hibernation for CachyOS with Limine bootloader
# REQUIRES: sudo interactive

set -euo pipefail

# Source common functions
source "$(dirname "$0")/../common.sh" || {
	echo "[ERROR] Could not source common.sh" >&2
	exit 1
}

readonly LIMINE_CONFIG="/etc/default/limine"
readonly MKINITCPIO_CONFIG="/etc/mkinitcpio.conf"

check_requirements() {
	local swap_uuid
	swap_uuid=$(blkid -t TYPE=swap -o value -s UUID | head -n1)

	if [[ -z "$swap_uuid" ]]; then
		log_error "No swap partition found"
		log_info "Please create a swap partition or swapfile before running this script"
		exit 1
	fi

	log_info "Found swap partition UUID: $swap_uuid"
	readonly SWAP_UUID="$swap_uuid"
}

escape_for_sed() {
	printf '%s\n' "$1" | sed 's/[[\.*^$()+?{|]/\\&/g'
}

show_diff_and_confirm() {
	local file="$1"
	local description="$2"
	local temp_file="$3"

	log_info "Proposed changes for $description"
	log_info "File: $file"
	echo

	if command_exists colordiff; then
		colordiff -u "$file" "$temp_file" || true
	else
		diff -u "$file" "$temp_file" || true
	fi

	echo
	confirm "Apply these changes?"
}

update_limine_config() {
	log_info "Updating Limine configuration"

	backup_file "$LIMINE_CONFIG"

	local current_cmdline new_cmdline changes_made=false
	current_cmdline=$(grep 'KERNEL_CMDLINE\[default\]' "$LIMINE_CONFIG" | cut -d'"' -f2)
	log_info "Current kernel cmdline: $current_cmdline"

	new_cmdline="$current_cmdline"

	if [[ ! "$new_cmdline" =~ resume= ]]; then
		new_cmdline="$new_cmdline resume=UUID=$SWAP_UUID"
		log_info "Adding resume parameter"
		changes_made=true
	fi

	if [[ ! "$new_cmdline" =~ ipv6.disable= ]]; then
		new_cmdline="$new_cmdline ipv6.disable=1"
		log_info "Adding ipv6.disable parameter"
		changes_made=true
	fi

	new_cmdline=$(echo "$new_cmdline" | sed 's/  */ /g' | sed 's/^ *//' | sed 's/ *$//')

	if [[ "$changes_made" == "true" ]]; then
		log_info "New kernel cmdline: $new_cmdline"

		local new_cmdline_escaped temp_file
		new_cmdline_escaped=$(escape_for_sed "$new_cmdline")
		temp_file=$(mktemp)

		sed "s|KERNEL_CMDLINE\[default\]=\".*\"|KERNEL_CMDLINE[default]=\"$new_cmdline_escaped\"|" "$LIMINE_CONFIG" >"$temp_file"

		if show_diff_and_confirm "$LIMINE_CONFIG" "LIMINE CONFIGURATION" "$temp_file"; then
			cp "$temp_file" "$LIMINE_CONFIG"
			log_info "Limine configuration updated successfully"
			rm "$temp_file"
			return 0
		else
			log_info "Skipping Limine configuration update"
			log_info "You can manually add: resume=UUID=$SWAP_UUID to your kernel cmdline"
			rm "$temp_file"
			return 1
		fi
	else
		log_info "Limine configuration already contains required parameters"
		return 1
	fi
}

update_mkinitcpio_config() {
	log_info "Updating mkinitcpio configuration"

	backup_file "$MKINITCPIO_CONFIG"

	if grep -q '^HOOKS=.*resume' "$MKINITCPIO_CONFIG"; then
		log_info "Resume hook already present in mkinitcpio.conf"
		return 1
	fi

	log_info "Adding resume hook to mkinitcpio.conf"

	local current_hooks new_hooks temp_file
	current_hooks=$(grep '^HOOKS=' "$MKINITCPIO_CONFIG")
	log_info "Current HOOKS: $current_hooks"

	new_hooks=$(echo "$current_hooks" | sed 's/filesystems/filesystems resume/' | sed 's/resume resume/resume/')

	if [[ ! "$new_hooks" =~ resume ]]; then
		new_hooks=$(echo "$current_hooks" | sed 's/)/ resume)/')
	fi

	log_info "New HOOKS: $new_hooks"

	local new_hooks_escaped
	new_hooks_escaped=$(escape_for_sed "$new_hooks")
	temp_file=$(mktemp)

	sed "s|^HOOKS=.*|$new_hooks_escaped|" "$MKINITCPIO_CONFIG" >"$temp_file"

	if show_diff_and_confirm "$MKINITCPIO_CONFIG" "MKINITCPIO CONFIGURATION" "$temp_file"; then
		cp "$temp_file" "$MKINITCPIO_CONFIG"
		log_info "mkinitcpio.conf updated successfully"
		rm "$temp_file"
		return 0
	else
		log_info "Skipping mkinitcpio.conf update"
		log_info "You can manually add 'resume' to your HOOKS array"
		rm "$temp_file"
		return 1
	fi
}

apply_system_changes() {
	local limine_updated="$1"
	local mkinitcpio_updated="$2"
	local operations_needed=()

	log_info "System updates required"

	if [[ "$mkinitcpio_updated" == "true" ]]; then
		operations_needed+=("Regenerate initramfs (mkinitcpio -P)")
	fi

	if [[ "$limine_updated" == "true" ]]; then
		operations_needed+=("Update Limine bootloader (limine-update)")
	fi

	if [[ ${#operations_needed[@]} -eq 0 ]]; then
		log_info "No system updates needed - configuration is already up to date"
		return 0
	fi

	log_info "The following operations need to be performed:"
	for i in "${!operations_needed[@]}"; do
		log_info "  $((i + 1)). ${operations_needed[$i]}"
	done

	if ! confirm "Apply all pending system changes?" "y"; then
		log_warn "Skipping system updates"
		log_warn "Configuration files have been updated but system changes were not applied"
		log_info "You will need to manually run:"

		if [[ "$mkinitcpio_updated" == "true" ]]; then
			log_info "  mkinitcpio -P"
		fi

		if [[ "$limine_updated" == "true" ]]; then
			log_info "  limine-update"
		fi

		log_warn "Hibernation will not work until these commands are executed"
		return 1
	fi

	log_info "Applying system changes"

	if [[ "$mkinitcpio_updated" == "true" ]]; then
		log_info "Regenerating initramfs"
		if ! mkinitcpio -P; then
			log_error "Failed to regenerate initramfs"
			log_info "Please check the configuration and run 'mkinitcpio -P' manually"
			exit 1
		fi
		log_info "Initramfs regenerated successfully"
	fi

	if [[ "$limine_updated" == "true" ]]; then
		log_info "Updating Limine configuration"
		if ! limine-update; then
			log_error "Failed to update Limine bootloader"
			log_info "Please run 'limine-update' manually"
			exit 1
		fi
		log_info "Limine bootloader updated successfully"
	fi

	return 0
}

show_completion_summary() {
	local limine_updated="$1"
	local mkinitcpio_updated="$2"

	log_info "Hibernation configuration completed successfully"
	echo
	log_info "Configuration details:"
	log_info "  - Swap UUID: $SWAP_UUID"

	if [[ "$limine_updated" == "true" ]]; then
		log_info "  - Resume parameter added to kernel cmdline"
		log_info "  - IPv6 disabled (optional optimization)"
	else
		log_info "  - Kernel cmdline already configured"
	fi

	if [[ "$mkinitcpio_updated" == "true" ]]; then
		log_info "  - Resume hook added to initramfs"
	else
		log_info "  - Resume hook already present in initramfs"
	fi

	echo
	log_info "Backups created:"
	log_info "  - $LIMINE_CONFIG.backup.*"
	log_info "  - $MKINITCPIO_CONFIG.backup.*"
	echo
	log_info "Next steps:"
	log_info "  1. Reboot your system"
	log_info "  2. Test hibernation with: systemctl hibernate"
	log_info "  3. Wake up and verify everything works correctly"
	echo
	log_info "If you encounter issues, restore backups with:"
	log_info "  cp $LIMINE_CONFIG.backup.* $LIMINE_CONFIG"
	log_info "  cp $MKINITCPIO_CONFIG.backup.* $MKINITCPIO_CONFIG"
	log_info "  mkinitcpio -P && limine-update"
}

main() {
	init_script

	check_requirements

	local limine_updated mkinitcpio_updated

	if update_limine_config; then
		limine_updated="true"
	else
		limine_updated="false"
	fi

	if update_mkinitcpio_config; then
		mkinitcpio_updated="true"
	else
		mkinitcpio_updated="false"
	fi

	if apply_system_changes "$limine_updated" "$mkinitcpio_updated"; then
		show_completion_summary "$limine_updated" "$mkinitcpio_updated"
		finish_script 0
	else
		finish_script 1
	fi
}

main "$@"