#!/bin/bash

# VS Code Emergency Memory Management
# Use this script when VS Code becomes unresponsive during agent iterations

echo "🔧 VS Code Emergency Memory Management"
echo "======================================"

# Function to show current status
show_status() {
    echo "📊 Current VS Code Status:"
    echo "Processes: $(ps aux | grep -E "(code|vscode)" | grep -v grep | wc -l)"
    echo "Memory: $(ps -o pid,ppid,pcpu,pmem,rss,args -C code | awk 'NR>1{sum+=$5} END{print sum/1024 " MB"}')"
    echo ""
}

# Function for gentle cleanup
gentle_cleanup() {
    echo "🧹 Performing gentle cleanup..."
    
    # Kill high-memory extension processes
    for pid in $(ps aux | grep 'extensions.*server' | awk '$6 > 300000 {print $2}'); do
        if [ -n "$pid" ]; then
            echo "Stopping high-memory extension process: $pid"
            kill -15 "$pid" 2>/dev/null
        fi
    done
    
    # Clean temporary files
    rm -rf /tmp/mcp-* 2>/dev/null
    find ~/.config/Code -name "*.log" -size +50M -delete 2>/dev/null
    
    sleep 3
    echo "✅ Gentle cleanup completed"
}

# Function for aggressive cleanup
aggressive_cleanup() {
    echo "⚠️  Performing aggressive cleanup..."
    
    # Stop all extension hosts
    pkill -f "extensionHost" 2>/dev/null
    
    # Stop language servers
    pkill -f "pylance.*server" 2>/dev/null
    pkill -f "language.*server" 2>/dev/null
    
    # Kill utility processes with high memory
    for pid in $(ps aux | grep 'code.*utility' | awk '$6 > 500000 {print $2}'); do
        if [ -n "$pid" ]; then
            echo "Force stopping utility process: $pid"
            kill -9 "$pid" 2>/dev/null
        fi
    done
    
    sleep 5
    echo "✅ Aggressive cleanup completed"
}

# Function to restart VS Code safely
restart_vscode() {
    echo "🔄 Restarting VS Code..."
    
    # Save current workspace
    WORKSPACE_PATH="$PWD"
    
    # Gracefully close VS Code
    pkill -15 code 2>/dev/null
    sleep 5
    
    # Force close if still running
    pkill -9 code 2>/dev/null
    sleep 2
    
    # Clean up leftover processes
    pkill -f "vscode" 2>/dev/null
    
    echo "VS Code stopped. Starting in 3 seconds..."
    sleep 3
    
    # Restart VS Code with the current workspace
    cd "$WORKSPACE_PATH"
    nohup code . > /dev/null 2>&1 &
    
    echo "✅ VS Code restarted"
}

# Interactive menu
case "${1:-menu}" in
    "status")
        show_status
        ;;
    "gentle")
        show_status
        gentle_cleanup
        show_status
        ;;
    "aggressive") 
        show_status
        aggressive_cleanup
        show_status
        ;;
    "restart")
        restart_vscode
        ;;
    "full")
        show_status
        echo "Performing full cleanup and restart..."
        gentle_cleanup
        sleep 2
        aggressive_cleanup
        sleep 2
        restart_vscode
        ;;
    "menu")
        show_status
        echo "Available actions:"
        echo "  status     - Show current VS Code status"
        echo "  gentle     - Light cleanup (kill high-memory extensions)"
        echo "  aggressive - Heavy cleanup (kill language servers, utilities)"
        echo "  restart    - Full VS Code restart"
        echo "  full       - Complete cleanup + restart"
        echo ""
        echo "Usage: $0 {status|gentle|aggressive|restart|full}"
        ;;
    *)
        echo "Unknown option: $1"
        $0 menu
        ;;
esac