#!/usr/bin/env bash

# Author: Santhosh Siva
# Date Created: 03-08-2025

# Description:
# This script creates a new git worktree for a specified branch.
# If the branch does not exist, it prompts the user to create it.

# 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
stash=false
target_branch=
custom_worktree_name=

set_flags() {
	while [ $# -gt 0 ]; do
		case "$1" in
		-h | --help)
			echo "g-wa - create git worktree for branch"
			echo " "
			echo "g-wa [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                             specify custom worktree directory name"
			echo "--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
			;;
		--worktree-name=*)
			custom_worktree_name="${1#*=}"
			if [ -z "$custom_worktree_name" ]; then
				echo ""
				echo "${RED}Error: No worktree name specified.$NC"
				exit 1
			fi
			;;
		--worktree-name)
			shift
			if [ $# -gt 0 ]; then
				custom_worktree_name="$1"
			else
				echo ""
				echo "${RED}Error: No worktree name specified.$NC"
				exit 1
			fi
			;;
		-s | --stash-changes)
			stash=true
			;;
		*)
			echo "${RED}Unknown option:${NC} $1"
			exit
			;;
		esac
		shift
	done
}

try_fetching_branch() {
	if ! fetch_changes "${target_branch}"; then
		print_message "" -1
		print_message "${RED}Failed to fetch changes for branch ${target_branch}. [Fail]${NC}" -1
		exit 1
	fi
}

check_branch_in_worktree() {
	local branch=$1
	local step_number=$2

	print_message "${BLUE}Checking if ${NC}${branch}${BLUE} branch exists...${NC}" $step_number

	# Check if branch is already checked out in another worktree
	local existing_worktree
	existing_worktree=$(git worktree list 2>/dev/null | grep -F "[${branch}]" | awk '{print $1}')
	if [ $? -ne 0 ]; then
		print_message "" -1
		print_message "${RED}Failed to check worktree list. [Fail]${NC}" -1
		exit 1
	fi
	if [ -n "$existing_worktree" ]; then
		print_message "${RED}Branch ${NC}${branch}${RED} is already checked out in worktree:${NC}"
		print_message "${GREEN}${existing_worktree}${NC}"
		print_message "${BLUE}Navigate to it with:${NC} ${GREEN}cd ${existing_worktree}${NC}"
		exit 0
	fi
}

generate_worktree_path() {
	local target_branch=$1
	local -n result_var=$2

	# Determine worktree directory based on current location
	local worktree_dir
	local current_dir=$(basename "$PWD")
	local parent_dir=$(basename "$(dirname "$PWD")")

	# Check if we're already in a worktree (parent directory contains worktrees)
	if [ "$parent_dir" = "worktrees" ]; then
		# We're in a worktree, use parent of parent for new worktrees
		worktree_dir="../"
	else
		# We're in main repo, create worktrees subdirectory
		if [ ! -d "../worktrees" ]; then
			mkdir ../worktrees
		fi
		worktree_dir="../worktrees/"
	fi

	# Use custom worktree name if provided, otherwise sanitize branch name
	local worktree_name
	if [ -n "$custom_worktree_name" ]; then
		worktree_name="$custom_worktree_name"
	else
		sanitize_branch_name "$target_branch" worktree_name
		if [ $? -ne 0 ]; then
			return 1
		fi
	fi
	local worktree_path="${worktree_dir}${worktree_name}"

	# Convert to absolute path
	local abs_path
	abs_path=$(cd "${worktree_path%/*}" 2>/dev/null && echo "$(pwd)/${worktree_path##*/}")
	local abs_path_exit=$?
	if [ $abs_path_exit -ne 0 ] || [ -z "$abs_path" ]; then
		print_message "" -1
		print_message "${RED}Failed to resolve absolute path. [Fail]${NC}" -1
		return 1
	fi
	worktree_path="$abs_path"

	# If path exists, generate a unique path
	if [ -d "$worktree_path" ]; then
		print_message "${PROMPT}Worktree path already exists.${NC}"
		local parent_dir="${worktree_path%/*}"
		local no_of_identical_dirs
		no_of_identical_dirs=$(find "$parent_dir" -maxdepth 1 -type d -name "${worktree_name}*" 2>/dev/null | wc -l)
		local find_exit=$?
		if [ $find_exit -ne 0 ]; then
			print_message "" -1
			print_message "${RED}Failed to count existing directories. [Fail]${NC}" -1
			return 1
		fi
		worktree_path="${parent_dir}/${worktree_name}_$(($no_of_identical_dirs + 1))"
		print_message "${BLUE}Using new path: ${NC}${worktree_path}${NC}"

		# Validate the new path
		if [ -z "$worktree_path" ]; then
			print_message "" -1
			print_message "${RED}Failed to generate unique worktree path. [Fail]${NC}" -1
			return 1
		fi
	fi

	# Set the result via nameref
	result_var="$worktree_path"
	return 0
}

checkout_or_create_branch() {
	local step_number=$1

	# Check if branch is already in a worktree
	check_branch_in_worktree "${target_branch}" $step_number

	# Generate worktree path
	local path
	generate_worktree_path "${target_branch}" path
	if [ $? -ne 0 ]; then
		print_message "" -1
		print_message "${RED}Failed to generate worktree path. [Fail]${NC}" -1
		exit 1
	fi

	# Check if branch exists locally
	if git show-ref --verify --quiet "refs/heads/${target_branch}"; then
		try_fetching_branch
		print_message "${BLUE}Creating ${NC}${target_branch}${BLUE} worktree...${NC}" $((step_number + 1))
		create_worktree "${target_branch}" "${path}"
		return
	fi

	print_message "${RED}Branch not found locally.${NC}"
	print_message "${BLUE}Checking if branch exists on remote...${NC}" $((step_number + 1))

	if git ls-remote --heads origin "${target_branch}" | grep -q "${target_branch}"; then
		print_message "${GREEN}Branch available on remote.${NC}"
		try_fetching_branch
		print_message "${BLUE}Creating ${NC}${target_branch}${BLUE} worktree...${NC}" $((step_number + 2))
		create_worktree "${target_branch}" "${path}"
		return
	fi

	print_message "${RED}Branch not found on remote.${NC}"

	# Get current branch and repo info for the "Will create" message
	local current_branch repo_name sanitized_target
	get_current_branch current_branch
	get_repo_name repo_name
	sanitize_branch_name "$target_branch" sanitized_target
	if [ $? -ne 0 ]; then
		print_message "" -1
		print_message "${RED}Failed to sanitize branch name. [Fail]${NC}" -1
		exit 1
	fi

	print_message "" -1
	print_message "${BLUE}Will create ${NC}${target_branch}${BLUE} worktree (cut from ${NC}${current_branch}${BLUE})${NC}" 0
	print_message "Location: ${BLUE}${repo_name}/worktrees/${sanitized_target}${NC}" 0
	print_message "" -1

	create_new_branch=$(prompt_user true "Create new branch?" $((step_number + 2)))
	if [ "${create_new_branch}" = "y" ]; then
		print_message "${BLUE}Creating ${NC}${target_branch}${BLUE} worktree...${NC}" $((step_number + 3))
		create_worktree "${target_branch}" "${path}" true
		return
	fi

	print_message "" -1
	print_message "${RED}Aborted.${NC}" -1
	exit 1
}

prompt_restructure_confirmation() {
	local git_root=$1
	local repo_name=$2
	local default_branch=$3
	local step_number=$4
	local target_branch=$5

	# Save the current branch before restructuring
	local original_branch
	get_current_branch original_branch
	print_message "${step_number}. ${BLUE}Repository Reorganization Required${NC}" -1
	print_message "Your repository will be moved into a 'main' subdirectory to enable worktree support" 0
	print_message "" 0
	print_message "Current location: ${BLUE}${git_root}${NC}" 0
	print_message "New location:     ${BLUE}${git_root}/main${NC}" 0
	print_message "Current branch:   ${BLUE}${original_branch}${NC}" 0
	print_message "" 0
	print_message "This structure allows you to work on multiple branches simultaneously" 0
	print_message "by creating worktrees in separate directories" 0

	# If not on default branch, mention worktree will be created
	if [ "$original_branch" != "$default_branch" ]; then
		local sanitized_original sanitized_target
		sanitize_branch_name "$original_branch" sanitized_original
		sanitize_branch_name "$target_branch" sanitized_target
		print_message "" 0
		print_message "Note: You're currently on ${BLUE}${original_branch}${NC} (not ${BLUE}${default_branch}${NC})" 0

		# Check if target branch is the same as original branch
		if [ "$target_branch" = "$original_branch" ]; then
			print_message "Two directories will be created:" 0
			print_message "" 0
			print_message "• ${BLUE}${repo_name}/main${NC} → ${BLUE}${default_branch}${NC} branch" 0
			print_message "• ${BLUE}${repo_name}/worktrees/${sanitized_original}${NC} → ${BLUE}${original_branch}${NC} branch" 0
		else
			print_message "Three directories will be created:" 0
			print_message "" 0
			print_message "• ${BLUE}${repo_name}/main${NC} → ${BLUE}${default_branch}${NC} branch" 0
			print_message "• ${BLUE}${repo_name}/worktrees/${sanitized_original}${NC} → ${BLUE}${original_branch}${NC} branch" 0
			print_message "• ${BLUE}${repo_name}/worktrees/${sanitized_target}${NC} → ${BLUE}${target_branch}${NC} branch (cut from ${BLUE}${original_branch}${NC})" 0
		fi
	fi

	# Prompt user for confirmation
	local proceed=$(prompt_user false "Proceed with restructuring?" "$step_number")
	if [ "$proceed" != "y" ]; then
		print_message "" -1
		print_message "${RED}Restructuring cancelled by user.${NC}" -1
		exit 1
	fi
}

create_main_dir() {
	local git_root=$1
	local temp_dir=$2
	local step_number=$3
	local -n result_var=$4

	# Create temp directory and move everything there
	if ! mkdir -p "$temp_dir"; then
		print_message "" -1
		print_message "${RED}Failed to create temporary directory. [Fail]${NC}" -1
		exit 1
	fi

	print_message "${BLUE}Moving repository contents to temporary location...${NC}" $step_number
	if ! (shopt -s dotglob && mv "$git_root"/* "$temp_dir/" 2>&1); then
		print_message "" -1
		print_message "${RED}Failed to move contents to temporary directory. [Fail]${NC}" -1
		rm -rf "$temp_dir"
		exit 1
	fi

	# Create main directory
	# Note: must NOT name this variable "main_dir" — that would shadow the nameref
	# result_var=$4 refers to a variable named "main_dir" in the caller's scope,
	# and bash resolves namerefs to the innermost local with that name.
	local _new_main_dir="${git_root}/main"
	if ! mkdir -p "$_new_main_dir"; then
		print_message "" -1
		print_message "${RED}Failed to create main directory. [Fail]${NC}" -1
		# Restore from temp
		(shopt -s dotglob && mv "$temp_dir"/* "$git_root/")
		rm -rf "$temp_dir"
		exit 1
	fi

	# Move everything from temp to main
	if ! (shopt -s dotglob && mv "$temp_dir"/* "$_new_main_dir/" 2>&1); then
		print_message "" -1
		print_message "${RED}Failed to move contents to main directory. [Fail]${NC}" -1
		exit 1
	fi

	# Clean up temp directory
	rm -rf "$temp_dir"

	result_var="$_new_main_dir"
	return 0
}

checkout_main_in_main_branch() {
	local repo_name=$1
	local main_dir=$2
	local default_branch=$3
	local step_number=$4

	# Checkout the default branch into main
	print_message "${BLUE}Creating main worktree...${NC}" $step_number
	git -c color.ui=always checkout "$default_branch" 2>&1 | indent
	if [ ${PIPESTATUS[0]} -ne 0 ]; then
		print_message "" -1
		print_message "${RED}Failed to checkout ${NC}${default_branch}${RED} branch. [Fail]${NC}" -1
		exit 1
	fi

	print_message "${GREEN}Repository restructured successfully to ${NC}${repo_name}/main/${GREEN}. [DONE]${NC}"
}

checkout_current_branch() {
	local step_number=$1
	local default_branch=$2
	local original_branch=$3

	if [ -z "$step_number" ]; then
		step_number=0
	fi

	if [ -z "$default_branch" ]; then
		print_message "" -1
		print_message "${RED}Default branch is not set. [Fail]${NC}" -1
		exit 1
	fi

	if [ "$default_branch" = "$original_branch" ]; then
		return 0
	fi

	local path
	generate_worktree_path "$original_branch" path
	if [ $? -ne 0 ]; then
		print_message "" -1
		print_message "${RED}Failed to generate worktree path. [Fail]${NC}" -1
		exit 1
	fi
	print_message "${BLUE}Creating ${NC}${original_branch}${BLUE} worktree...${NC}" $step_number
	create_worktree "$original_branch" "$path"

	# Navigate to the created worktree
	if ! navigate_to_dir "$path"; then
		print_message "" -1
		print_message "${RED}Failed to navigate to ${original_branch} worktree. [Fail]${NC}" -1
		exit 1
	fi
}

already_restructured() {
	local current_dir=$(pwd)
	if ! navigate_to_main_dir "false"; then
		return 1
	fi
	if ! navigate_to_dir "$current_dir"; then
		print_message "" -1
		print_message "${RED}Failed to navigate back to original directory. [Fail]${NC}" -1
		exit 1
	fi
	return 0
}

create_main_dir_if_necessary() {
	local default_branch=$1
	local step_number=$2

	if [ -z "$step_number" ]; then
		step_number=0
	fi

	local git_root current_dir parent_dir repo_name
	if ! get_repo_info git_root current_dir parent_dir repo_name; then
		print_message "" -1
		print_message "${RED}Failed to get repository information. [Fail]${NC}" -1
		exit 1
	fi

	if already_restructured; then
		return 0
	fi

	# Save original branch before restructuring
	local original_branch
	get_current_branch original_branch

	# Prompt user for restructuring
	prompt_restructure_confirmation "$git_root" "$repo_name" "$default_branch" "$step_number" "$target_branch"

	# Get parent directory path
	local parent_path="${git_root%/*}"
	local temp_dir="${parent_path}/.temp-${repo_name}-$$"

	# Create main directory and restructure
	local main_dir
	create_main_dir "$git_root" "$temp_dir" $((step_number + 1)) main_dir

	# Navigate to main directory using direct path
	if ! navigate_to_dir "$main_dir"; then
		print_message "" -1
		print_message "${RED}Failed to navigate to main directory. [Fail]${NC}" -1
		exit 1
	fi

	# Checkout the default branch into main directory
	checkout_main_in_main_branch "$repo_name" "$main_dir" "$default_branch" $((step_number + 2))

	# Create worktree for original branch and navigate to it (if different from default)
	checkout_current_branch $((step_number + 3)) "$default_branch" "$original_branch"

	return 1
}

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

	check_if_target_branch_is_set $target_branch

	if ! is_git_repo "$PWD"; then
		print_message "" -1
		print_message "${RED}Not in a git repository. [Fail]${NC}" -1
		exit 1
	fi

	already_on_branch $target_branch 1

	stash_changes $stash 2

	if has_uncommitted_changes "$PWD"; then
		print_message "" -1
		print_message "${RED}Uncommitted changes detected. Commit or stash changes before proceeding. [Fail]${NC}" -1
		exit 1
	fi

	local default_branch current_branch
	get_default_branch default_branch
	get_current_branch current_branch

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

	# Check if we need to restructure (to calculate proper step numbers)
	local git_root current_dir parent_dir repo_name
	if ! get_repo_info git_root current_dir parent_dir repo_name; then
		print_message "" -1
		print_message "${RED}Failed to get repository information. [Fail]${NC}" -1
		exit 1
	fi
	create_main_dir_if_necessary "$default_branch" $step_number
	local needs_restructure=$?

	if [ $needs_restructure -eq 1 ]; then
		checkout_or_create_branch $((step_number + 4))
	else
		checkout_or_create_branch $step_number
	fi
}

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