#!/bin/bash

# Environment Cleanup Script
# Safely removes conda environments and virtual environments

set -e

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

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

log() { echo -e "${GREEN}[INFO]${NC} $1"; }
warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
info() { echo -e "${BLUE}[INFO]${NC} $1"; }

# List all environments
list_environments() {
    echo "======================================"
    echo "Available Environments"
    echo "======================================"
    
    echo -e "\n${BLUE}Conda Environments:${NC}"
    if command -v conda &> /dev/null; then
        conda env list | grep -v "^#" | grep -v "^$" || echo "  No conda environments found"
    else
        echo "  Conda not available"
    fi
    
    echo -e "\n${BLUE}Virtual Environments in current directory:${NC}"
    for dir in "${SCRIPT_DIR}"/*/; do
        if [ -f "${dir}bin/activate" ]; then
            echo "  - $(basename "$dir")"
        fi
    done | sort || echo "  No virtual environments found"
    
    echo "======================================"
}

# Remove conda environment
remove_conda_env() {
    local env_name="$1"
    
    if ! conda env list | grep -q "^${env_name} "; then
        error "Conda environment '${env_name}' not found"
        return 1
    fi
    
    warning "About to remove conda environment: ${env_name}"
    read -p "Are you sure? (y/n): " -n 1 -r
    echo
    
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log "Removing conda environment: ${env_name}"
        conda env remove -n "${env_name}" -y
        log "✓ Conda environment removed: ${env_name}"
    else
        info "Cancelled"
    fi
}

# Remove virtual environment
remove_venv() {
    local venv_path="$1"
    
    if [ ! -d "$venv_path" ]; then
        # Check if it's a relative path
        venv_path="${SCRIPT_DIR}/${venv_path}"
        if [ ! -d "$venv_path" ]; then
            error "Virtual environment not found: $1"
            return 1
        fi
    fi
    
    if [ ! -f "${venv_path}/bin/activate" ]; then
        error "Not a valid virtual environment: ${venv_path}"
        return 1
    fi
    
    warning "About to remove virtual environment: ${venv_path}"
    echo "This will delete the entire directory and all its contents"
    read -p "Are you sure? (y/n): " -n 1 -r
    echo
    
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log "Removing virtual environment: ${venv_path}"
        rm -rf "${venv_path}"
        log "✓ Virtual environment removed"
    else
        info "Cancelled"
    fi
}

# Clean build artifacts
clean_build_artifacts() {
    log "Cleaning build artifacts..."
    
    # Clean Python cache
    find "${SCRIPT_DIR}" -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true
    find "${SCRIPT_DIR}" -type f -name "*.pyc" -delete 2>/dev/null || true
    find "${SCRIPT_DIR}" -type f -name "*.pyo" -delete 2>/dev/null || true
    
    # Clean build directories
    for dir in build dist *.egg-info; do
        if [ -d "${SCRIPT_DIR}/${dir}" ]; then
            log "Removing ${dir}/"
            rm -rf "${SCRIPT_DIR}/${dir}"
        fi
    done
    
    # Clean Flash Attention build if exists
    if [ -d "${SCRIPT_DIR}/flash-attention/build" ]; then
        log "Cleaning Flash Attention build artifacts"
        rm -rf "${SCRIPT_DIR}/flash-attention/build"
        rm -rf "${SCRIPT_DIR}/flash-attention/dist"
        rm -rf "${SCRIPT_DIR}/flash-attention/*.egg-info"
    fi
    
    log "✓ Build artifacts cleaned"
}

# Clean logs
clean_logs() {
    log "Cleaning log files..."
    
    local log_count=$(find "${SCRIPT_DIR}" -name "*.log" 2>/dev/null | wc -l)
    
    if [ "$log_count" -eq 0 ]; then
        info "No log files to clean"
        return
    fi
    
    echo "Found $log_count log files"
    read -p "Remove all log files? (y/n): " -n 1 -r
    echo
    
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        find "${SCRIPT_DIR}" -name "*.log" -delete
        log "✓ Log files cleaned"
    else
        info "Keeping log files"
    fi
}

# Clean temporary files
clean_temp_files() {
    log "Cleaning temporary files..."
    
    # Clean pip cache for current user
    if [ -d ~/.cache/pip ]; then
        read -p "Clean pip cache? This may slow down future installations (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf ~/.cache/pip/*
            log "✓ Pip cache cleaned"
        fi
    fi
    
    # Clean conda package cache
    if command -v conda &> /dev/null; then
        read -p "Clean conda package cache? This may slow down future installations (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            conda clean -a -y > /dev/null 2>&1
            log "✓ Conda cache cleaned"
        fi
    fi
}

# Interactive cleanup menu
interactive_cleanup() {
    while true; do
        echo ""
        echo "======================================"
        echo "Environment Cleanup Menu"
        echo "======================================"
        echo "1. List all environments"
        echo "2. Remove conda environment"
        echo "3. Remove virtual environment"
        echo "4. Clean build artifacts"
        echo "5. Clean log files"
        echo "6. Clean temporary/cache files"
        echo "7. Full cleanup (artifacts + logs + temp)"
        echo "0. Exit"
        echo "======================================"
        
        read -p "Select option: " choice
        
        case $choice in
            1)
                list_environments
                ;;
            2)
                list_environments
                read -p "Enter conda environment name to remove: " env_name
                remove_conda_env "$env_name"
                ;;
            3)
                echo "Virtual environments in current directory:"
                for dir in "${SCRIPT_DIR}"/*/; do
                    if [ -f "${dir}bin/activate" ]; then
                        echo "  - $(basename "$dir")"
                    fi
                done
                read -p "Enter virtual environment name/path to remove: " venv_path
                remove_venv "$venv_path"
                ;;
            4)
                clean_build_artifacts
                ;;
            5)
                clean_logs
                ;;
            6)
                clean_temp_files
                ;;
            7)
                log "Running full cleanup..."
                clean_build_artifacts
                clean_logs
                clean_temp_files
                log "✓ Full cleanup complete"
                ;;
            0)
                echo "Exiting..."
                exit 0
                ;;
            *)
                error "Invalid option"
                ;;
        esac
    done
}

# Main execution
main() {
    # If arguments provided, process them
    if [ $# -gt 0 ]; then
        case "$1" in
            --list|-l)
                list_environments
                ;;
            --conda|-c)
                if [ -z "$2" ]; then
                    error "Please provide conda environment name"
                fi
                remove_conda_env "$2"
                ;;
            --venv|-v)
                if [ -z "$2" ]; then
                    error "Please provide virtual environment path"
                fi
                remove_venv "$2"
                ;;
            --artifacts)
                clean_build_artifacts
                ;;
            --logs)
                clean_logs
                ;;
            --all)
                clean_build_artifacts
                clean_logs
                clean_temp_files
                ;;
            --help|-h)
                echo "Usage: $0 [options]"
                echo "Options:"
                echo "  --list, -l           List all environments"
                echo "  --conda, -c NAME     Remove conda environment"
                echo "  --venv, -v PATH      Remove virtual environment"
                echo "  --artifacts          Clean build artifacts"
                echo "  --logs               Clean log files"
                echo "  --all                Clean everything (artifacts, logs, temp)"
                echo "  --help, -h           Show this help"
                echo ""
                echo "Run without arguments for interactive mode"
                ;;
            *)
                error "Unknown option: $1. Use --help for usage."
                ;;
        esac
    else
        # Interactive mode
        interactive_cleanup
    fi
}

main "$@"