#!/bin/bash
# Cleanup git worktrees
# Usage: worktree-cleanup.sh [feature-name]
# If no feature-name provided, lists candidates for cleanup

set -e

FEATURE_NAME="$1"

# Get project root from current working directory (where Claude Code is running)
PROJECT_ROOT="$(pwd)"
PROJECT_NAME=$(basename "$PROJECT_ROOT")
WORKTREES_DIR="$(dirname "$PROJECT_ROOT")/ab-worktrees"

echo "========================================"
echo "  Worktree Cleanup"
echo "========================================"
echo ""

cd "$PROJECT_ROOT"

# If specific feature provided, cleanup that one
if [ -n "$FEATURE_NAME" ]; then
    BRANCH_NAME="feature/ab-${FEATURE_NAME}"
    WORKTREE_PATH="${WORKTREES_DIR}/${PROJECT_NAME}-${FEATURE_NAME}"

    echo "Cleaning up: $FEATURE_NAME"
    echo "  Branch: $BRANCH_NAME"
    echo "  Path:   $WORKTREE_PATH"
    echo ""

    # Check if worktree exists
    if [ ! -d "$WORKTREE_PATH" ]; then
        echo "Worktree not found at: $WORKTREE_PATH"
        exit 1
    fi

    # Check if branch is merged
    MAIN_BRANCH=$(git symbolic-ref refs/remotes/origin/HEAD 2>/dev/null | sed 's@^refs/remotes/origin/@@' || echo "main")

    if git branch --merged "$MAIN_BRANCH" | grep -q "$BRANCH_NAME"; then
        echo "Branch is merged to $MAIN_BRANCH"
    else
        echo "Warning: Branch is NOT merged to $MAIN_BRANCH"
        echo "Proceed with cleanup? (y/n)"
        read -r CONFIRM
        if [ "$CONFIRM" != "y" ]; then
            echo "Aborted"
            exit 0
        fi
    fi

    # Remove worktree
    echo "Removing worktree..."
    git worktree remove "$WORKTREE_PATH" --force

    # Remove branch
    echo "Removing branch..."
    git branch -D "$BRANCH_NAME" 2>/dev/null || echo "Branch already removed or doesn't exist"

    # Update registry
    REGISTRY_FILE="$PROJECT_ROOT/.auto-build-data/worktrees/worktree-registry.json"
    if [ -f "$REGISTRY_FILE" ] && command -v jq &> /dev/null; then
        TEMP_FILE=$(mktemp)
        jq --arg name "$FEATURE_NAME" '.worktrees |= map(select(.name != $name))' "$REGISTRY_FILE" > "$TEMP_FILE" && mv "$TEMP_FILE" "$REGISTRY_FILE"
        echo "Registry updated"
    fi

    echo ""
    echo "Cleanup complete!"

else
    # List all worktrees and show which can be cleaned up
    echo "Checking worktrees for cleanup..."
    echo ""

    MAIN_BRANCH=$(git symbolic-ref refs/remotes/origin/HEAD 2>/dev/null | sed 's@^refs/remotes/origin/@@' || echo "main")

    echo "Merged branches (safe to cleanup):"
    MERGED_BRANCHES=$(git branch --merged "$MAIN_BRANCH" | grep "feature/ab-" || echo "")

    if [ -n "$MERGED_BRANCHES" ]; then
        echo "$MERGED_BRANCHES" | while read -r branch; do
            FEATURE=$(echo "$branch" | sed 's/.*feature\/ab-//')
            echo "  - $FEATURE"
        done
    else
        echo "  (none)"
    fi

    echo ""
    echo "Unmerged branches:"
    UNMERGED_BRANCHES=$(git branch --no-merged "$MAIN_BRANCH" | grep "feature/ab-" || echo "")

    if [ -n "$UNMERGED_BRANCHES" ]; then
        echo "$UNMERGED_BRANCHES" | while read -r branch; do
            FEATURE=$(echo "$branch" | sed 's/.*feature\/ab-//')
            echo "  - $FEATURE (NOT merged)"
        done
    else
        echo "  (none)"
    fi

    echo ""
    echo "To cleanup a specific worktree:"
    echo "  /ab:worktree cleanup <feature-name>"
fi