#!/bin/bash
set -euo pipefail

# ── Colors ───────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# ── Logging helpers ──────────────────────────────────────────────────
log_info()  { printf "${CYAN}[INFO]${NC}  %s\n" "$*"; }
log_warn()  { printf "${YELLOW}[WARN]${NC}  %s\n" "$*"; }
log_ok()    { printf "${GREEN}[ OK ]${NC}  %s\n" "$*"; }
log_error() { printf "${RED}[ERR]${NC}   %s\n" "$*" >&2; }

# ── Usage ────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
Usage: ${0##*/} [OPTIONS]

Update all git repositories found under the current directory.

Options:
    -h, --help          Show this help
    -m, --main          Checkout to main/master branch after update
    -i, --ignore=DIR    Ignore folder(s) (repeatable, comma-separated)
    -o, --only=DIR      Update only these folder(s) (repeatable, comma-separated)

Examples:
    ${0##*/} --ignore=vendor,node_modules
    ${0##*/} --only=project-a,project-b --main
EOF
    exit 1
}

# ── Parse arguments ──────────────────────────────────────────────────
IGNORE_FOLDERS=()
ONLY_FOLDERS=()
FLAG_SWITCH_TO_MAIN=""

# Map long options to short ones
translated=()
while (( $# )); do
    case "$1" in
        --help)      translated+=("-h"); shift ;;
        --main)      translated+=("-m"); shift ;;
        --ignore=*)  translated+=("-i" "${1#*=}"); shift ;;
        --only=*)    translated+=("-o" "${1#*=}"); shift ;;
        --)          translated+=("--"); shift; translated+=("$@"); break ;;
        *)           translated+=("$1"); shift ;;
    esac
done
set -- "${translated[@]}"

while getopts ":hmi:o:" opt; do
    case "$opt" in
        h)  usage ;;
        m)  FLAG_SWITCH_TO_MAIN=1 ;;
        i)
            # Support comma-separated values
            IFS=',' read -ra parts <<< "$OPTARG"
            IGNORE_FOLDERS+=("${parts[@]}")
            ;;
        o)
            IFS=',' read -ra parts <<< "$OPTARG"
            ONLY_FOLDERS+=("${parts[@]}")
            ;;
        :)  log_error "Missing argument for -${OPTARG}"; usage ;;
        \?) log_error "Invalid option: -${OPTARG}"; usage ;;
    esac
done
shift $((OPTIND - 1))

# ── Debug: show filters ─────────────────────────────────────────────
if (( ${#IGNORE_FOLDERS[@]} > 0 )); then
    log_info "Ignoring folders: ${IGNORE_FOLDERS[*]}"
fi
if (( ${#ONLY_FOLDERS[@]} > 0 )); then
    log_info "Only updating folders: ${ONLY_FOLDERS[*]}"
fi

# ── Find git repositories ───────────────────────────────────────────
find_repos() {
    local find_args=(. -name ".git" -type d)

    # Whitelist: only matching paths
    if (( ${#ONLY_FOLDERS[@]} > 0 )); then
        local only_paths=()
        for folder in "${ONLY_FOLDERS[@]}"; do
            only_paths+=("-o" "-path" "*/${folder}/*")
        done
        # Remove the leading -o and wrap in parentheses
        unset 'only_paths[0]'
        find_args+=("(" "${only_paths[@]}" ")")
    fi

    # Blacklist: excluded paths
    for folder in "${IGNORE_FOLDERS[@]}"; do
        find_args+=("-not" "-path" "*/${folder}/*")
    done

    find "${find_args[@]}"
}

# ── Update a single repository ───────────────────────────────────────
update_repo() {
    local repo_dir="$1"
    local had_stash=0

    # Stash uncommitted changes if any
    if [[ -n "$(git -C "$repo_dir" status --porcelain)" ]]; then
        log_warn "Uncommitted changes in ${repo_dir} — stashing"
        git -C "$repo_dir" stash --quiet
        had_stash=1
    fi

    # Determine main and current branches
    local master_branch current_branch
    master_branch="$(git -C "$repo_dir" symbolic-ref refs/remotes/origin/HEAD 2>/dev/null \
                     | sed 's@^refs/remotes/origin/@@')" || true
    current_branch="$(git -C "$repo_dir" branch --show-current)"

    if [[ -z "$master_branch" ]]; then
        log_warn "Cannot determine default branch for ${repo_dir} — skipping"
        return
    fi

    if [[ "$master_branch" == "$current_branch" ]]; then
        git -C "$repo_dir" pull --all --quiet
    else
        git -C "$repo_dir" checkout "$master_branch" --quiet
        git -C "$repo_dir" pull --all --quiet
        if [[ -z "$FLAG_SWITCH_TO_MAIN" ]]; then
            git -C "$repo_dir" checkout "$current_branch" --quiet 2>/dev/null \
                || log_warn "Branch '${current_branch}' was deleted on remote"
        fi
    fi

    # Prune stale remote-tracking branches
    git -C "$repo_dir" remote prune origin --quiet 2>/dev/null || true

    # Delete local branches that have been merged into the default branch
    local merged
    merged="$(git -C "$repo_dir" branch --merged \
              | grep -v -E "^\*|${master_branch}" || true)"
    if [[ -n "$merged" ]]; then
        echo "$merged" | xargs git -C "$repo_dir" branch -D --quiet 2>/dev/null || true
    fi

    # Restore stashed changes
    if (( had_stash )); then
        git -C "$repo_dir" stash pop --quiet || log_warn "Stash pop failed in ${repo_dir}"
    fi

    log_ok "$repo_dir"
}

# ── Main ─────────────────────────────────────────────────────────────
main() {
    local count=0

    while IFS= read -r git_dir; do
        local repo_dir
        repo_dir="$(dirname "$git_dir")"
        log_info "Updating ${repo_dir} …"
        update_repo "$repo_dir"
        (( count++ )) || true
    done < <(find_repos)

    if (( count == 0 )); then
        log_warn "No git repositories found"
    else
        log_ok "Done — updated ${count} repo(s)"
    fi
}

main