#!/bin/bash

# Setup Git Worktrees for Parallel Development
# This script creates separate worktrees for different development streams

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
WORKTREE_BASE="$(dirname "$PROJECT_ROOT")/trax-worktrees"

echo "🌳 Setting up Git Worktrees for Trax"
echo "=================================="
echo "Project Root: $PROJECT_ROOT"
echo "Worktree Base: $WORKTREE_BASE"
echo ""

# Create worktree base directory
mkdir -p "$WORKTREE_BASE"

# Function to create a worktree
create_worktree() {
    local name=$1
    local branch=$2
    local description=$3
    local worktree_path="$WORKTREE_BASE/$name"
    
    echo "📁 Creating worktree: $name"
    echo "   Branch: $branch"
    echo "   Path: $worktree_path"
    echo "   Purpose: $description"
    
    # Check if branch exists remotely
    if git ls-remote --heads origin "$branch" | grep -q "$branch"; then
        echo "   ✓ Branch exists remotely, checking out..."
        git worktree add "$worktree_path" "origin/$branch"
    else
        echo "   → Creating new branch..."
        git worktree add -b "$branch" "$worktree_path"
    fi
    
    # Setup virtual environment for the worktree
    echo "   🐍 Setting up virtual environment..."
    cd "$worktree_path"
    python3.11 -m venv .venv
    source .venv/bin/activate
    pip install --quiet --upgrade pip
    pip install --quiet uv
    uv pip install -e ".[dev]" --quiet
    deactivate
    
    # Create .env.local if it doesn't exist
    if [ ! -f "$worktree_path/.env.local" ]; then
        echo "# Local environment overrides" > "$worktree_path/.env.local"
    fi
    
    # Create a README for the worktree
    cat > "$worktree_path/WORKTREE_README.md" << EOF
# Worktree: $name

**Branch**: $branch
**Purpose**: $description

## Quick Start

\`\`\`bash
# Activate virtual environment
source .venv/bin/activate

# Run tests
uv run pytest

# Start development
# ... your commands here ...
\`\`\`

## Switching Between Worktrees

\`\`\`bash
# List all worktrees
git worktree list

# Switch to another worktree
cd $WORKTREE_BASE/<worktree-name>
\`\`\`
EOF
    
    echo "   ✅ Worktree created successfully!"
    echo ""
}

# Main execution
cd "$PROJECT_ROOT"

# Ensure we're on main branch and up to date
echo "🔄 Updating main branch..."
git checkout main
git pull origin main 2>/dev/null || echo "   (No remote changes)"
echo ""

# Create worktrees for different development streams
echo "🚀 Creating development worktrees..."
echo ""

# 1. Feature Development
create_worktree "trax-features" "feature/development" \
    "New feature development and experimentation"

# 2. Testing & QA
create_worktree "trax-testing" "testing/qa" \
    "Testing, QA, and validation work"

# 3. Documentation
create_worktree "trax-docs" "docs/updates" \
    "Documentation updates and improvements"

# 4. Performance Optimization
create_worktree "trax-performance" "perf/optimization" \
    "Performance tuning and optimization"

# 5. Bug Fixes
create_worktree "trax-bugfix" "fix/current" \
    "Bug fixes and hotfixes"

# Create convenience script for switching
cat > "$WORKTREE_BASE/switch.sh" << 'EOF'
#!/bin/bash
# Quick switcher for Trax worktrees

WORKTREE_BASE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

echo "🌳 Trax Worktrees:"
echo ""

# List worktrees with numbers
worktrees=($(ls -d $WORKTREE_BASE/trax-* 2>/dev/null | xargs -n1 basename))
for i in "${!worktrees[@]}"; do
    branch=$(cd "$WORKTREE_BASE/${worktrees[$i]}" && git branch --show-current)
    echo "  $((i+1)). ${worktrees[$i]} [$branch]"
done

echo ""
read -p "Select worktree (1-${#worktrees[@]}): " choice

if [[ $choice -ge 1 && $choice -le ${#worktrees[@]} ]]; then
    selected="${worktrees[$((choice-1))]}"
    echo "Switching to $selected..."
    cd "$WORKTREE_BASE/$selected"
    exec $SHELL
else
    echo "Invalid choice"
fi
EOF

chmod +x "$WORKTREE_BASE/switch.sh"

# Create status script
cat > "$WORKTREE_BASE/status.sh" << 'EOF'
#!/bin/bash
# Show status of all Trax worktrees

WORKTREE_BASE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

echo "🌳 Trax Worktree Status"
echo "======================="
echo ""

for worktree in $WORKTREE_BASE/trax-*/; do
    if [ -d "$worktree" ]; then
        name=$(basename "$worktree")
        cd "$worktree"
        branch=$(git branch --show-current)
        status=$(git status --porcelain | wc -l | xargs)
        ahead_behind=$(git status -sb | head -1 | grep -oE '\[.*\]' || echo "[synced]")
        
        echo "📁 $name"
        echo "   Branch: $branch $ahead_behind"
        if [ "$status" -gt 0 ]; then
            echo "   Changes: $status uncommitted files"
        else
            echo "   Status: Clean"
        fi
        echo ""
    fi
done

echo "---"
echo "Run '$WORKTREE_BASE/switch.sh' to switch between worktrees"
EOF

chmod +x "$WORKTREE_BASE/status.sh"

# Summary
echo ""
echo "✅ Worktree Setup Complete!"
echo "=========================="
echo ""
echo "📁 Worktrees created in: $WORKTREE_BASE"
echo ""
echo "🔧 Available worktrees:"
git worktree list | sed 's/^/   /'
echo ""
echo "📝 Convenience scripts:"
echo "   • $WORKTREE_BASE/switch.sh - Switch between worktrees"
echo "   • $WORKTREE_BASE/status.sh - Show status of all worktrees"
echo ""
echo "💡 Tips:"
echo "   • Each worktree has its own .venv and can run independently"
echo "   • Use 'git worktree list' to see all worktrees"
echo "   • Use 'git worktree remove <path>' to remove a worktree"
echo "   • Open multiple Claude Code sessions - one per worktree"
echo ""
echo "🚀 To start developing:"
echo "   cd $WORKTREE_BASE/<worktree-name>"
echo "   source .venv/bin/activate"
echo "   claude  # Start Claude Code"