#!/usr/bin/env bash

# Author: Santhosh Siva
# Date Created: 12-02-2025

# Description:
# This script removes a git worktree for a specified branch.
# It searches through the worktrees directory and removes matching worktrees.

# Resolve script directory (works with symlinks for npm global install)
SOURCE="${BASH_SOURCE[0]}"
while [ -L "$SOURCE" ]; do
  DIR="$(cd -P "$(dirname "$SOURCE")" && pwd)"
  SOURCE="$(readlink "$SOURCE")"
  [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE"
done
SCRIPT_DIR="$(cd -P "$(dirname "$SOURCE")" && pwd)"
source "$SCRIPT_DIR/utils"

# Default Values
target_branch=
worktree_name=
stash=false

set_flags() {
	while [ $# -gt 0 ]; do
		case "$1" in
		-h | --help)
			echo "g-wr - remove a git worktree for a specified branch or name"
			echo " "
			echo "g-wr [options]"
			echo " "
			echo "options:"
			echo "-h, --help                                                             show brief help"
			echo "--target-branch BRANCH, --target-branch=BRANCH, -t=BRANCH, -t BRANCH   specify the target branch"
			echo "--worktree-name NAME, --worktree-name=NAME, -w=NAME, -w NAME           specify the worktree directory name"
			echo "-s, --stash-changes                                                    stash changes before proceeding"
			exit 0
			;;
		-t=* | --target-branch=*)
			target_branch="${1#*=}"
			if [ -z "$target_branch" ]; then
				echo ""
				echo "${RED}Error: No target branch specified.$NC"
				exit 1
			fi
			;;
		-t | --target-branch)
			shift
			if [ $# -gt 0 ]; then
				target_branch="$1"
			else
				echo ""
				echo "${RED}Error: No target branch specified.$NC"
				exit 1
			fi
			;;
		-w=* | --worktree-name=*)
			worktree_name="${1#*=}"
			if [ -z "$worktree_name" ]; then
				echo ""
				echo "${RED}Error: No worktree name specified.$NC"
				exit 1
			fi
			;;
		-w | --worktree-name)
			shift
			if [ $# -gt 0 ]; then
				worktree_name="$1"
			else
				echo ""
				echo "${RED}Error: No worktree name specified.$NC"
				exit 1
			fi
			;;
		-s | --stash-changes)
			stash=true
			;;
		*)
			echo ""
			echo "${RED}Unknown option:${NC} $1"
			exit 1
			;;
		esac
		shift
	done
}

cd_to_worktrees_dir() {
	# Get git root and navigate to repo root's worktrees directory
	local git_root
	git_root=$(git rev-parse --show-toplevel 2>&1)
	if [ $? -ne 0 ] || [ -z "$git_root" ]; then
		print_message "" -1
		print_message "${RED}Failed to get git root directory. [Fail]${NC}" -1
		exit 1
	fi
	local git_root_parent=$(dirname "$git_root")
	local git_root_parent_name=$(basename "$git_root_parent")

	# If we're in a worktree (parent is "worktrees"), use that directory
	if [ "$git_root_parent_name" = "worktrees" ]; then
		if ! navigate_to_dir "$git_root_parent"; then
			print_message "" -1
			print_message "${RED}Failed to navigate to worktrees directory. [Fail]${NC}" -1
			exit 1
		fi
		return 0
	fi

	# Otherwise, we're in main, so look for ../worktrees
	local worktrees_dir="$git_root_parent/worktrees"
	if [ ! -d "$worktrees_dir" ]; then
		print_message "" -1
		print_message "${RED}Worktrees directory not found at ${NC}${worktrees_dir}${RED}. [Fail]${NC}" -1
		exit 1
	fi

	if ! navigate_to_dir "$worktrees_dir"; then
		print_message "" -1
		print_message "${RED}Failed to navigate to worktrees directory. [Fail]${NC}" -1
		exit 1
	fi
}

find_worktree_by_name() {
	local worktree_name=$1

	cd_to_worktrees_dir
	local worktrees_dir=$(pwd)

	for dir in *; do
		if [ ! -d "$dir" ]; then
			continue
		fi

		if [ "$dir" = "$worktree_name" ]; then
			echo "$worktrees_dir/$dir"
			return 0
		fi
	done

	return 1
}

find_worktree_by_branch() {
	local target_branch=$1

	check_if_target_branch_is_set "$target_branch"

	cd_to_worktrees_dir
	local worktrees_dir=$(pwd)

	for dir in *; do
		if [ ! -d "$dir" ]; then
			continue
		fi

		local worktree_path="$worktrees_dir/$dir"

		# Use subshell to check branch without changing parent directory
		local branch
		branch=$(cd "$worktree_path" 2>/dev/null && git rev-parse --abbrev-ref HEAD 2>/dev/null)
		local branch_exit=$?
		if [ $branch_exit -ne 0 ] || [ -z "$branch" ]; then
			continue
		fi

		if [ "$branch" = "$target_branch" ]; then
			echo "$worktree_path"
			return 0
		fi
	done

	return 1
}

find_dir_to_remove() {
	if [ -n "$worktree_name" ]; then
		find_worktree_by_name "$worktree_name"
		return $?
	else
		find_worktree_by_branch "$target_branch"
		return $?
	fi
}

check_if_worktree_is_clean() {
	local worktree_path=$1
	local step_number=$2

	print_message "${BLUE}Checking for uncommitted changes...${NC}" $step_number

	if has_uncommitted_changes "$worktree_path"; then
		print_message "${PROMPT}Worktree has uncommitted changes.${NC}" 0
		response=$(prompt_user "false" "Do you want to discard them and remove the worktree?" $step_number)
		if [ "$response" = "n" ]; then
			print_message "" -1
			print_message "${RED}Worktree removal cancelled. Your changes are safe.${NC}" -1
			exit 0
		fi
	else
		print_message "${GREEN}No uncommitted changes found.${NC}" 0
	fi
}

remove_worktree() {
	local worktree_path=$1
	local step_number=$2

	print_message "${BLUE}Removing worktree at ${NC}${worktree_path}${BLUE}...${NC}" $step_number

	local worktree_parent
	worktree_parent=$(dirname "$worktree_path")

	# Always try using git worktree remove first with --force --force to handle
	# both uncommitted changes and untracked files
	if git -c color.ui=always worktree remove --force --force "$worktree_path" 2>&1 | indent; then
		# Verify the directory was actually removed
		if [ ! -d "$worktree_path" ]; then
			# Navigate to parent in case we were running from inside the removed worktree
			cd "$worktree_parent" 2>/dev/null
			print_message "${GREEN}Worktree removed successfully.${NC}" 0
			return 0
		fi
	fi

	# Fallback to rm -rf if git worktree remove failed
	print_message "${PROMPT}Git worktree remove incomplete. Attempting manual removal...${NC}" 0
	if rm -rf "$worktree_path" 2>&1 | indent; then
		# Navigate to parent before pruning since the worktree directory no longer exists
		cd "$worktree_parent" 2>/dev/null
		# Also need to prune the worktree from git's records
		git -c color.ui=always worktree prune 2>&1 | indent
		# Verify the directory was actually removed
		if [ ! -d "$worktree_path" ]; then
			print_message "${GREEN}Worktree removed successfully.${NC}" 0
			return 0
		fi
	fi

	print_message "${RED}Failed to remove worktree. [Fail]${NC}" 0
	return 1
}

main() {
	set_flags "$@"
	validate_dependencies git figlet lolcat
	print_banner

	local step_number=1
	if [ "$stash" = true ]; then
		stash_changes $stash $step_number
		step_number=2
	fi

	if [ -n "$worktree_name" ]; then
		print_message "${BLUE}Searching for worktree with name ${NC}${worktree_name}${BLUE}...${NC}" $step_number
	else
		print_message "${BLUE}Searching for worktree with branch ${NC}${target_branch}${BLUE}...${NC}" $step_number
	fi

	local worktree_path
	worktree_path=$(find_dir_to_remove)
	if [ $? -ne 0 ] || [ -z "$worktree_path" ]; then
		print_message "" -1
		if [ -n "$worktree_name" ]; then
			print_message "${RED}No worktree found with name ${NC}${worktree_name}${RED}. [Fail]${NC}" -1
		else
			print_message "${RED}No worktree found for branch ${NC}${target_branch}${RED}. [Fail]${NC}" -1
		fi
		exit 1
	fi

	worktree_dirname=$(basename "$worktree_path")
	if [ -n "$worktree_name" ]; then
		print_message "${GREEN}Found worktree: ${NC}${worktree_dirname}${NC}" $((step_number + 1))
	else
		print_message "${GREEN}Found worktree: ${NC}${worktree_dirname}${GREEN} at ${NC}${worktree_path}${NC}" $((step_number + 1))
	fi

	# Prevent removal of main directory
	if [ "$worktree_dirname" = "main" ]; then
		print_message "" -1
		print_message "${RED}Cannot remove main directory. This is your default branch worktree. [Fail]${NC}" -1
		exit 1
	fi

	check_if_worktree_is_clean "$worktree_path" $((step_number + 2))

	remove_worktree "$worktree_path" $((step_number + 3))
}


# Only run main if script is executed directly (not sourced)
if [ "${BASH_SOURCE[0]}" = "${0}" ]; then
	main "$@"
	exit 0
fi
