ai-digest

#!/bin/bash

# Script: ai-digest
# Description: Compiles an AI digest from multiple folders/files and copies it to clipboard
# Author: Cooper Morgan
# Date: 2024-11-21

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

# Function to print status messages
print_status() {
    echo -e "${CYAN}=>${NC} $1"
}

# Function to print success messages
print_success() {
    echo -e "${GREEN}✓${NC} $1"
}

# Function to print error messages
print_error() {
    echo -e "${RED}✗${NC} $1" >&2
}

# Function to print debug messages
print_debug() {
    [ "$DEBUG" = true ] && echo -e "${YELLOW}debug:${NC} $1" >&2
}

# Default values
PATHS=()                                         # Array to store multiple input paths
IGNORE_PATHS=("node_modules" ".*" "__pycache__" ".git" "package-lock.json", "build", "assets") # Auto-ignores
TEMP_FILE="/tmp/ai-digest-$$.md"
CLIPBOARD_CMD=""
MAX_SIZE_MB=10 # Maximum size limit in MB
CURRENT_SIZE=0
DEBUG=false

# Check for required commands
if ! command -v bc >/dev/null 2>&1; then
    print_error "The 'bc' command is required but not found."
    echo -e "${BLUE}For Debian/Ubuntu, install bc with:${NC}"
    echo -e "    sudo apt-get install bc"
    echo -e "${BLUE}For macOS, install bc with:${NC}"
    echo -e "    brew install bc"
    exit 1
fi

# Detect operating system and set clipboard command
if [[ "$OSTYPE" == "darwin"* ]]; then
    CLIPBOARD_CMD="pbcopy"
elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
    if command -v xclip >/dev/null 2>&1; then
        CLIPBOARD_CMD="xclip -selection clipboard"
    elif command -v xsel >/dev/null 2>&1; then
        CLIPBOARD_CMD="xsel --clipboard --input"
    else
        print_error "No clipboard command found."
        echo -e "${BLUE}For Debian/Ubuntu, install xclip with:${NC}"
        echo -e "    sudo apt-get install xclip"
        echo -e "${BLUE}Or install xsel with:${NC}"
        echo -e "    sudo apt-get install xsel"
        exit 1
    fi
fi

# Function to show help message
show_help() {
    echo -e "
    ${GREEN}AI Digest Generator${NC}
    A tool to create markdown digests of files and directories

    ${BOLD}USAGE${NC}
      $(basename "$0") [paths...] [options]

    ${BOLD}OPTIONS${NC}
      ${CYAN}-i${NC}, ${CYAN}--ignore${NC}   PATTERN   Additional paths to ignore (repeatable)
      ${CYAN}-m${NC}, ${CYAN}--max-size${NC} SIZE      Maximum size in MB (default: 10)
      ${CYAN}-d${NC}, ${CYAN}--debug${NC}              Enable debug output
      ${CYAN}-h${NC}, ${CYAN}--help${NC}               Show this help message

    ${BOLD}EXAMPLES${NC}
      $(basename "$0")                            # Process current directory
      $(basename "$0") path/to/file               # Process single file
      $(basename "$0") dir1/ dir2/                # Process multiple directories
      $(basename "$0") src/ tests/ -i vendor      # Process directories, ignore vendor
      $(basename "$0") * -i node_modules          # Process all files in current directory

    ${BOLD}NOTES${NC}
      • When processing directories, node_modules and dotfiles (.*) are ignored by default
      • The output is automatically copied to your clipboard
      • Maximum file size can be adjusted with --max-size (in MB)
      • Use --debug to see detailed processing information"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
    -i | --ignore)
        IGNORE_PATHS+=("$2")
        shift 2
        ;;
    -m | --max-size)
        MAX_SIZE_MB="$2"
        shift 2
        ;;
    -d | --debug)
        DEBUG=true
        shift
        ;;
    -h | --help)
        show_help
        exit 0
        ;;
    -*)
        print_error "Unknown option: $1"
        show_help
        exit 1
        ;;
    *)
        # Convert to absolute path
        if [ -e "$1" ]; then
            if [ -f "$1" ]; then
                PATHS+=("$(cd "$(dirname "$1")" && pwd)/$(basename "$1")")
            elif [ -d "$1" ]; then
                PATHS+=("$(cd "$1" && pwd)")
            fi
        else
            print_error "Path does not exist: $1"
            exit 1
        fi
        shift
        ;;
    esac
done

# If no paths were provided, use current directory
if [ ${#PATHS[@]} -eq 0 ]; then
    print_debug "No paths provided, using current directory"
    PATHS+=("$(pwd)")
fi

# Check if clipboard command is available
if [ -z "$CLIPBOARD_CMD" ]; then
    print_error "No clipboard command found. Please install xclip or xsel on Linux, or use macOS."
    exit 1
fi

# Debug info
if [ "$DEBUG" = true ]; then
    print_debug "Working directory: $(pwd)"
    print_debug "Input paths: ${PATHS[*]}"
    print_debug "Ignore patterns: ${IGNORE_PATHS[*]}"
fi

# Function to check if path should be ignored
should_ignore() {
    local check_path="$1"
    local base_name=$(basename "$check_path")

    # Check against ignore patterns
    for ignore_path in "${IGNORE_PATHS[@]}"; do
        # Handle glob patterns (like .*)
        if [[ "$base_name" == $ignore_path ]]; then
            return 0
        fi
        # Handle exact matches
        if [[ "$check_path" == *"$ignore_path"* ]]; then
            return 0
        fi
    done
    return 1
}

# Function to check file size and update total
check_size() {
    local new_size=$(stat -f%z "$TEMP_FILE" 2>/dev/null || stat -c%s "$TEMP_FILE")
    CURRENT_SIZE=$((new_size / 1024 / 1024)) # Convert to MB

    if [ $CURRENT_SIZE -gt $MAX_SIZE_MB ]; then
        print_error "Maximum size limit (${MAX_SIZE_MB}MB) exceeded! Stopping..."
        return 1
    fi
    return 0
}

# Function to process a single file
process_file() {
    local file="$1"
    local base_path="$2"
    local relative_path="${file#$base_path}"
    relative_path="${relative_path#/}" # Remove leading slash

    # Skip if file should be ignored
    if should_ignore "$file"; then
        print_debug "Skipping ignored file: $file"
        return 0
    fi

    if file "$file" | grep -q text; then
        # Check file size before processing
        local file_size=$(stat -f%z "$file" 2>/dev/null || stat -c%s "$file")
        file_size=$((file_size / 1024 / 1024)) # Convert to MB

        if [ $file_size -gt $((MAX_SIZE_MB / 2)) ]; then
            print_debug "Skipping large file ($file_size MB): $file"
            return 0
        fi

        print_debug "Processing file: $file"

        # Add file information with just the relative path
        echo "## $relative_path" >>"$TEMP_FILE"

        # Determine file type for syntax highlighting
        local ext="${file##*.}"
        # If the file has no extension or is the same as the filename (no dot)
        if [ "$ext" = "$file" ] || [ -z "$ext" ]; then
            # Try to detect shell scripts
            if grep -q "^#\!/.*sh" "$file" || file "$file" | grep -qi "shell script"; then
                ext="bash"
            else
                # Default to txt if we can't determine the type
                ext="txt"
            fi
        fi
        printf '\n```%s\n' "$ext" >>"$TEMP_FILE"
        cat "$file" >>"$TEMP_FILE" || {
            print_error "Failed to read file: $file"
        }
        printf '\n```\n\n' >>"$TEMP_FILE"

        if ! check_size; then
            print_error "Size limit exceeded after processing: $file"
            return 1
        fi
    else
        print_debug "Skipping non-text file: $file"
    fi
    return 0
}

# Function to create markdown from directory structure
create_markdown() {
    local current_path="$1"
    local base_path="$2"

    # Skip if path should be ignored
    if should_ignore "$current_path"; then
        print_debug "Skipping ignored path: $current_path"
        return 0
    fi

    local any_files_processed=false
    
    # Build prune conditions for find
    local prune_conditions=""
    for ignore_path in "${IGNORE_PATHS[@]}"; do
        if [[ "$ignore_path" != ".*" ]]; then  # Handle dotfiles separately
            prune_conditions="$prune_conditions -name $ignore_path -o"
        fi
    done
    
    # Remove trailing "-o" if it exists
    prune_conditions="${prune_conditions% -o}"
    
    # Construct find command with pruning
    local find_cmd="find \"$current_path\""
    if [ ! -z "$prune_conditions" ]; then
        find_cmd="$find_cmd \( $prune_conditions \) -prune -o"
    fi
    # Add condition to ignore dotfiles
    find_cmd="$find_cmd -not -path '*/\.*'"
    # Finally, add the -type f to only find files
    find_cmd="$find_cmd -type f -print0"
    
    print_debug "Executing find command: $find_cmd"
    
    # Execute the constructed find command
    while IFS= read -r -d '' file; do
        if process_file "$file" "$base_path"; then
            any_files_processed=true
        else
            return 1
        fi
    done < <(eval "$find_cmd" | sort -z)

    [ "$any_files_processed" = false ] && print_debug "No files processed in: $current_path"
    return 0
}

# Function to format file size in human-readable format
format_size() {
    local size=$1
    local units=("B" "KB" "MB" "GB")
    local unit=0
    local max_unit=3 # units array length - 1

    while [ $unit -lt $max_unit ]; do
        # Use bc for decimal comparison
        if [ $(echo "$size < 1024" | bc) -eq 1 ]; then
            break
        fi
        size=$(echo "scale=2; $size/1024" | bc)
        unit=$((unit + 1))
    done

    # Round to 2 decimal places and trim trailing zeros
    size=$(echo $size | awk '{printf "%.2f\n", $0}' | sed 's/\.00$//')
    echo "${size}${units[$unit]}"
}

# Create temporary file
>"$TEMP_FILE"

# Add header to digest
echo "# AI Digest" >>"$TEMP_FILE"
echo "Generated on $(date '+%Y-%m-%d %H:%M:%S')" >>"$TEMP_FILE"
echo >>"$TEMP_FILE"

# Process each input path
for path in "${PATHS[@]}"; do
    print_status "Processing: $path"
    
    if [ -f "$path" ]; then
        # Process single file
        if ! process_file "$path" "$(dirname "$path")"; then
            print_error "Failed to process file: $path"
            rm "$TEMP_FILE"
            exit 1
        fi
    else
        # Add tree visualization for directory if available
        if command -v tree >/dev/null 2>&1; then
            echo "## Directory Structure: $(basename "$path")" >>"$TEMP_FILE"
            echo >>"$TEMP_FILE"
            echo '```' >>"$TEMP_FILE"

            # Combine all ignore patterns for tree
            all_ignores=()
            for pattern in "${IGNORE_PATHS[@]}"; do
                if [[ "$pattern" != ".*" ]]; then # Handle dotfiles separately
                    all_ignores+=("$pattern")
                fi
            done
            
            # Join patterns with |
            ignore_pattern=$(IFS='|'; echo "${all_ignores[*]}")
            
            # Run tree with combined ignore pattern
            tree --charset=ascii -I "$ignore_pattern" --dirsfirst "$path" 2>/dev/null >>"$TEMP_FILE"

            echo '```' >>"$TEMP_FILE"
            echo >>"$TEMP_FILE"
            echo "## Contents" >>"$TEMP_FILE"
            echo >>"$TEMP_FILE"
        fi

        # Generate markdown content for directory
        if ! create_markdown "$path" "$path"; then
            print_error "Failed to generate digest for: $path"
            rm "$TEMP_FILE"
            exit 1
        fi
    fi
done

# Check if we have any content
if [ ! -s "$TEMP_FILE" ]; then
    print_error "No content was generated"
    rm "$TEMP_FILE"
    exit 1
fi

# Copy to clipboard
if cat "$TEMP_FILE" | eval "$CLIPBOARD_CMD"; then
    print_success "Digest has been copied to clipboard!"
else
    print_error "Failed to copy to clipboard!"
    print_status "Digest has been saved to: $TEMP_FILE"
    exit 1
fi

# Get and display the actual file size
size_in_bytes=$(stat -f%z "$TEMP_FILE" 2>/dev/null || stat -c%s "$TEMP_FILE")
formatted_size=$(format_size $size_in_bytes)
print_status "Digest size: $formatted_size"

# Cleanup
rm "$TEMP_FILE"