#!/usr/bin/env bash
# LICENSE: unlicense. This is free and unencumbered software released into the public domain.
# see unlicense.org for full license
#
# source this command if you want to have it change directories for you

RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[0;33m"
CLEAR="\033[0m"
VERBOSE=

function usage {
    cat <<EOF
worktree [-v] <branch name>

create a git worktree with <branch name>. Will create a worktree if one isn't
found that matches the given name.

To create a worktree from a github fork, use the format username:branchname:
    worktree someuser:feature-branch

Will copy over some untracked files to the new worktree. By default, this includes
.env, .envrc, .env.local, .tool-versions, and mise.toml files.

To customize the list of untracked files to copy for a particular repository:
    git config --add worktree.untrackedfiles ".env"
    git config --add worktree.untrackedfiles "mise.toml"

To set a global configuration for all repositories:
    git config --global --add worktree.untrackedfiles ".env"
    git config --global --add worktree.untrackedfiles "mise.toml"

If you have any custom configuration set, it will override the defaults
completely, so add all files you want copied.

To configure large directories to copy in the background (defaults: node_modules):
    git config --add worktree.backgroundcopy ".venv"
    git config --add worktree.backgroundcopy "node_modules"

Background copy configuration also supports global settings with --global.
EOF
    kill -INT $$
}

function die {
    printf '%b%s%b\n' "$RED" "$1" "$CLEAR"
    # exit the script, but if it was sourced, don't kill the shell
    kill -INT $$
}

function warn {
    printf '%b%s%b\n' "$YELLOW" "$1" "$CLEAR"
}

# If at all possible, use copy-on-write to copy files. This is especially
# important to allow us to copy node_modules directories efficiently
#
# On mac or bsd: try to use -c
# see:
# https://spin.atomicobject.com/2021/02/23/git-worktrees-untracked-files/
#
# On gnu: use --reflink
#
# Use /bin/cp directly to avoid any of the user's aliases - this script is
# often eval'ed
#
# I tried to figure out how to actually determine the filesystem support for
# copy-on-write, but did not find any good references, so I'm falling back on
# "try and see if it fails"
function cp_cow {
    if ! /bin/cp -Rc "$1" "$2"; then
        if ! /bin/cp -R --reflink "$1" "$2"; then
            if ! /bin/cp -R "$1" "$2"; then
                warn "Unable to copy file $1 to $2 - folder may not exist"
            fi
        fi
    fi
}

# Get the list of untracked files to copy from git config. If no config is
# found, return the default list
function get_untracked_files_pattern {
    # Default file patterns
    local default_patterns="\.env|\.envrc|\.env\.local|\.mise\.toml|\.tool-versions|mise\.toml|Cargo\.lock"

    # Check if we have custom patterns in git config
    if git config --get-all worktree.untrackedfiles > /dev/null 2>&1; then
        local custom_patterns=""
        custom_patterns=$(git config --get-all worktree.untrackedfiles | tr '\n' '|' | sed 's/|$//')
        echo "^($custom_patterns)$"
    else
        echo "^($default_patterns)$"
    fi
}

# Get the list of directories to copy in the background from git config.
# If no config is found, return the default list (node_modules)
function get_background_copy_dirs {
    if git config --get-all worktree.backgroundcopy > /dev/null 2>&1; then
        git config --get-all worktree.backgroundcopy
    else
        echo "node_modules"
    fi
}

# Create a worktree from a given branchname, in exactly the way I like it.
function _worktree {
    if [ -z "$1" ]; then
        usage
    fi

    if [ -n "$VERBOSE" ]; then
        set -x
    fi
    branchname="$1"

    # if the branch is of the form user:branch, assume that is a fork 'branch'
    # of the current repository at github.com/user/<repo_name>
    if [[ "$branchname" == *:* ]]; then
        username="${branchname%%:*}"
        branchname="${branchname#*:}"
        remote="$username"

        # Check if remote exists
        if ! git remote get-url "$remote" > /dev/null 2>&1; then
            # Get the origin URL and construct the fork URL
            origin_url=$(git remote get-url origin)

            # Handle both SSH and HTTPS URLs
            if [[ "$origin_url" == git@github.com:* ]]; then
                # SSH format: git@github.com:owner/repo.git
                repo_name="${origin_url##*/}"
                fork_url="git@github.com:${username}/${repo_name}"
            elif [[ "$origin_url" == https://github.com/* ]]; then
                # HTTPS format: https://github.com/owner/repo.git
                repo_name="${origin_url##*/}"
                fork_url="https://github.com/${username}/${repo_name}"
            else
                die "Unable to determine fork URL from origin: $origin_url"
            fi

            warn "Adding remote '$remote' for $fork_url"
            if ! git remote add "$remote" "$fork_url"; then
                die "Failed to add remote '$remote'"
            fi
        fi

        # Fetch from the remote
        warn "Fetching from remote '$remote'"
        if ! git fetch "$remote"; then
            die "Failed to fetch from remote '$remote'"
        fi

        # Update the dirname to include the username prefix, and replace
        # slashes with underscores
        dirname="${username}_${branchname//\//_}"
    else
        remote="origin"
        # Replace slashes with underscores. If there's no slash, dirname will equal
        # branchname. So "alu/something-other" becomes "alu_something-other", but
        # "quick-fix" stays unchanged
        # https://www.tldp.org/LDP/abs/html/parameter-substitution.html
        dirname=${branchname//\//_}
    fi

    # Find the git repository root and change to it; if you run this command
    # from a directory within the repository this will make sure it still does
    # the right thing
    git_root=$(git rev-parse --show-toplevel)
    if [ -z "$git_root" ]; then
        die "Not in a git repository"
    fi

    # Change to the git root directory
    cd "$git_root" || die "Failed to change to git root directory"

    # pull the most recent version of the remote
    if ! git pull; then
        warn "Unable to run git pull, there may not be an upstream"
    fi

    # if the branch name already exists, we want to check it out. Otherwise,
    # create a new branch. I'm sure there's probably a way to do that in one
    # command, but I'm done fiddling with git at this point
    #
    # As far as I can tell, we have to check locally and remotely separately if
    # we want to be accurate. See https://stackoverflow.com/a/75040377 for the
    # reasoning here. Also this has some caveats, but probably works well
    # enough :shrug:
    #
    # if the branch exists locally:
    if git for-each-ref --format='%(refname:lstrip=2)' refs/heads | grep -E "^$branchname$" > /dev/null 2>&1; then
        if ! git worktree add "../$dirname" "$branchname"; then
            die "failed to create git worktree $branchname"
        fi
    # if the branch exists on a remote:
    elif git for-each-ref --format='%(refname:lstrip=3)' "refs/remotes/$remote" | grep -E "^$branchname$" > /dev/null 2>&1; then
        # create a local worktree tracking the remote branch
        if ! git worktree add -b "$branchname" "../$dirname" "$remote/$branchname"; then
            die "failed to create git worktree $remote/$branchname"
        fi
    else
        # otherwise, create a new branch
        if [ "$remote" != "origin" ]; then
            warn "Branch $remote/$branchname not found, creating new local branch"
        fi
        if ! git worktree add -b "$branchname" "../$dirname"; then
            die "failed to create git worktree $branchname"
        fi
    fi

    # Find untracked files that we want to copy to the new worktree
    #
    # Copy configured directories in the background These can contain large
    # dependency trees (node_modules, .venv, etc.) We don't want to copy nested
    # instances, only the root directory
    #
    # Because this can be so slow, we launch it as a process in the background,
    # and we can't use the cp_cow function, so use a similar command
    #
    # I tried running `npm i` after copying the packages, but it doesn't seem
    # to work for reasons that are unclear to me.
    while IFS= read -r dir; do
        if [ -d "$dir" ]; then
            tmpfile=$(mktemp)
            cp_cmd="cp -Rc $dir ../$dirname/$dir ||
                cp -R --reflink $dir ../$dirname/$dir"

            warn "copying $dir in the background"
            warn "logging to $tmpfile"
            nohup bash -c "($cp_cmd)" &> "$tmpfile" &
        fi
    done < <(get_background_copy_dirs)

    # get a list of untracked files to copy, either from "git config" or a
    # default list
    untracked_file_re=$(get_untracked_files_pattern)

    # this will fail for any files with \n in their names. don't do that.
    IFS=$'\n'

    # The rule we're using here is to copy untracked files if:
    # - They match the untracked file regex
    # - *and* they are in a tracked directory
    #
    # This means we'll copy Cargo.lock from the root, for example, but not from
    # inside a .gitignored `target` directory
    #
    # We ignore the background copy directories because we already handled them
    # above

    # Build exclusion arguments for fd
    fd_excludes=()
    while IFS= read -r dir; do
        fd_excludes+=(-E "$dir")
    done < <(get_background_copy_dirs)

    files_to_copy=()
    while IFS= read -r file; do
        parent_dir=$(dirname "$file")
        # Check if parent directory is tracked (or is root which is always tracked)
        if [ "$parent_dir" = "." ] || ! git check-ignore -q "$parent_dir" 2>/dev/null; then
            files_to_copy+=("$file")
        fi
    done < <(fd -u "$untracked_file_re" "${fd_excludes[@]}")

    for f in "${files_to_copy[@]}"; do
        cp_cow "$f" ../"$dirname"/"$f"
    done

    # return the shell to normal splitting mode
    unset IFS

    # if there was an envrc file, tell direnv that it's ok to run it
    if [ -f "../$dirname/.envrc" ]; then
        direnv allow "../$dirname"
    fi

    # now change to the new tree
    cd "../$dirname" || return

    # Initialize and update submodules if they exist
    if [ -f "$git_root/.gitmodules" ]; then
        warn "initializing submodules"
        if ! git submodule update --init --recursive; then
            warn "Failed to initialize submodules"
        fi
    fi

    printf "%bcreated worktree %s%b\n" "$GREEN" "../$dirname" "$CLEAR"
}

while true; do
    case $1 in
        help | -h | --help)
            usage
            ;;
        -v | --verbose)
            VERBOSE=true
            shift
            ;;
        *)
            break
            ;;
    esac
done

_worktree "$@"
