#!/bin/bash
#
# cleanup_worktrees.sh - Remove agent worktrees and tmux sessions
#
# This script cleans up resources created by create_worktrees.sh:
# - Kills tmux sessions
# - Removes git worktrees
# - Optionally deletes agent branches
#
# Usage:
#   ./cleanup_worktrees.sh [options]
#
# Options:
#   -p, --project PATH   Project root directory (default: current directory)
#   -a, --agent ID       Clean up specific agent only
#   -b, --branches       Also delete agent branches
#   -f, --force          Force removal even with uncommitted changes
#   -n, --dry-run        Show what would be done without executing
#   -h, --help           Show this help message
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# Default values
PROJECT_ROOT="$(pwd)"
SPECIFIC_AGENT=""
DELETE_BRANCHES=false
FORCE=false
DRY_RUN=false

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--project)
            PROJECT_ROOT="$2"
            shift 2
            ;;
        -a|--agent)
            SPECIFIC_AGENT="$2"
            shift 2
            ;;
        -b|--branches)
            DELETE_BRANCHES=true
            shift
            ;;
        -f|--force)
            FORCE=true
            shift
            ;;
        -n|--dry-run)
            DRY_RUN=true
            shift
            ;;
        -h|--help)
            head -30 "$0" | tail -25
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            exit 1
            ;;
    esac
done

# Functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[OK]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

get_agent_worktrees() {
    cd "$PROJECT_ROOT"
    git worktree list --porcelain | grep -E "^worktree.*agent-" | sed 's/worktree //' || true
}

get_agent_sessions() {
    tmux list-sessions -F "#{session_name}" 2>/dev/null | grep -E "^(claude-code|gemini|codex)" || true
}

cleanup_tmux_session() {
    local session_name=$1

    if ! tmux has-session -t "$session_name" 2>/dev/null; then
        log_info "tmux session '$session_name' not found"
        return 0
    fi

    log_info "Killing tmux session: $session_name"

    if [ "$DRY_RUN" = true ]; then
        echo "  Would run: tmux kill-session -t $session_name"
    else
        tmux kill-session -t "$session_name"
        log_success "Killed tmux session: $session_name"
    fi
}

cleanup_worktree() {
    local worktree_path=$1

    if [ ! -d "$worktree_path" ]; then
        log_info "Worktree not found: $worktree_path"
        return 0
    fi

    # Get branch name before removal
    local branch=""
    if [ -d "$worktree_path/.git" ] || [ -f "$worktree_path/.git" ]; then
        branch=$(cd "$worktree_path" && git branch --show-current 2>/dev/null || true)
    fi

    log_info "Removing worktree: $worktree_path"

    if [ "$DRY_RUN" = true ]; then
        if [ "$FORCE" = true ]; then
            echo "  Would run: git worktree remove $worktree_path --force"
        else
            echo "  Would run: git worktree remove $worktree_path"
        fi
    else
        cd "$PROJECT_ROOT"
        if [ "$FORCE" = true ]; then
            git worktree remove "$worktree_path" --force
        else
            git worktree remove "$worktree_path"
        fi
        log_success "Removed worktree: $worktree_path"
    fi

    # Optionally delete branch
    if [ "$DELETE_BRANCHES" = true ] && [ -n "$branch" ]; then
        cleanup_branch "$branch"
    fi
}

cleanup_branch() {
    local branch=$1

    # Don't delete main/master branches
    if [[ "$branch" == "main" || "$branch" == "master" ]]; then
        log_warn "Skipping protected branch: $branch"
        return 0
    fi

    log_info "Deleting branch: $branch"

    if [ "$DRY_RUN" = true ]; then
        if [ "$FORCE" = true ]; then
            echo "  Would run: git branch -D $branch"
        else
            echo "  Would run: git branch -d $branch"
        fi
    else
        cd "$PROJECT_ROOT"
        if [ "$FORCE" = true ]; then
            git branch -D "$branch" 2>/dev/null || log_warn "Could not delete branch: $branch"
        else
            git branch -d "$branch" 2>/dev/null || log_warn "Could not delete branch: $branch (may have unmerged changes)"
        fi
    fi
}

cleanup_specific_agent() {
    local agent_id=$1
    local worktree_path="$PROJECT_ROOT/../agent-$agent_id"

    echo ""
    log_info "Cleaning up agent: $agent_id"
    echo "============================================"

    # Kill tmux session
    cleanup_tmux_session "$agent_id"

    # Remove worktree
    cleanup_worktree "$worktree_path"

    echo "============================================"
    log_success "Agent $agent_id cleaned up"
}

cleanup_all_agents() {
    echo ""
    log_info "Finding all agent resources..."
    echo ""

    # Get all agent worktrees
    local worktrees
    worktrees=$(get_agent_worktrees)

    # Get all agent tmux sessions
    local sessions
    sessions=$(get_agent_sessions)

    # Also check for common agent directories
    for agent_id in "claude-code" "gemini" "codex"; do
        local worktree_path="$PROJECT_ROOT/../agent-$agent_id"

        # Check tmux session
        if tmux has-session -t "$agent_id" 2>/dev/null; then
            cleanup_tmux_session "$agent_id"
        fi

        # Check worktree
        if [ -d "$worktree_path" ]; then
            cleanup_worktree "$worktree_path"
        fi
    done

    # Clean up any additional worktrees found
    if [ -n "$worktrees" ]; then
        echo "$worktrees" | while read -r worktree_path; do
            if [ -n "$worktree_path" ]; then
                cleanup_worktree "$worktree_path"
            fi
        done
    fi
}

show_status() {
    echo ""
    log_info "Current status after cleanup:"
    echo "============================================"

    echo ""
    echo "Git Worktrees:"
    cd "$PROJECT_ROOT"
    git worktree list

    echo ""
    echo "tmux Sessions:"
    tmux list-sessions 2>/dev/null || echo "  No tmux sessions"

    echo "============================================"
}

# Main execution
main() {
    echo ""
    echo "========================================"
    echo "  Multi-Agent Workspace Cleanup"
    echo "========================================"
    echo ""

    if [ "$DRY_RUN" = true ]; then
        log_warn "DRY RUN MODE - No changes will be made"
        echo ""
    fi

    if [ "$FORCE" = true ]; then
        log_warn "FORCE MODE - Will force removal even with uncommitted changes"
        echo ""
    fi

    if [ "$DELETE_BRANCHES" = true ]; then
        log_warn "BRANCHES MODE - Will also delete agent branches"
        echo ""
    fi

    # Check if project is a git repo
    if [ ! -d "$PROJECT_ROOT/.git" ]; then
        log_error "$PROJECT_ROOT is not a git repository"
        exit 1
    fi

    if [ -n "$SPECIFIC_AGENT" ]; then
        cleanup_specific_agent "$SPECIFIC_AGENT"
    else
        cleanup_all_agents
    fi

    # Prune worktree references
    if [ "$DRY_RUN" = false ]; then
        log_info "Pruning worktree references..."
        cd "$PROJECT_ROOT"
        git worktree prune
    fi

    # Show final status
    show_status

    echo ""
    log_success "Cleanup complete!"
    echo ""
}

# Run main
main
