#!/bin/bash
# Working Interactive System Optimizer

set -euo pipefail

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

# Configuration
BACKUP_DIR="/var/lib/system-tuning/backups"
LOG_FILE="/var/log/system-optimization.log"

log() {
    echo -e "${BLUE}[INFO]${NC} $1"
    [[ -w "$(dirname "$LOG_FILE")" ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG_FILE"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
    [[ -w "$(dirname "$LOG_FILE")" ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] SUCCESS: $1" >> "$LOG_FILE"
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
    [[ -w "$(dirname "$LOG_FILE")" ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] WARNING: $1" >> "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
    [[ -w "$(dirname "$LOG_FILE")" ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1" >> "$LOG_FILE"
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root"
        echo "Usage: sudo $0"
        exit 1
    fi
}

analyze_and_prompt() {
    log "Analyzing current system optimizations..."
    
    # Get system info
    local ram_gb=$(free -g | awk '/^Mem:|^Speicher:/{print $2}')
    local cpu_cores=$(nproc)
    
    echo ""
    echo "📊 System Information:"
    echo "  💾 RAM: ${ram_gb}GB"
    echo "  🖥️  CPU: ${cpu_cores} cores"
    echo ""
    
    # Check and prompt for each optimization
    local needs_changes=false
    
    # === ZRAM CHECK ===
    echo "🗜️ zram Compressed Swap Analysis:"
    local optimal_size=$((ram_gb / 2))
    [[ $optimal_size -lt 2 ]] && optimal_size=2
    
    if [[ -e /dev/zram0 ]] && swapon --show | grep -q zram0; then
        local zram_size=$(swapon --show | grep zram0 | awk '{print $3}')
        local zram_gb=$(echo "$zram_size" | sed 's/[^0-9.]//g' | cut -d. -f1)
        
        echo "  📈 Current: $zram_size"
        echo "  📊 Optimal: ${optimal_size}GB"
        
        if [[ $zram_gb -gt $((optimal_size + 2)) ]] || [[ $zram_gb -lt $((optimal_size - 1)) ]]; then
            echo "  ⚠️  Current size is suboptimal"
            needs_changes=true
            read -p "  Would you like to reconfigure zram to ${optimal_size}GB? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                SETUP_ZRAM=true
            else
                SETUP_ZRAM=false
            fi
        else
            echo "  ✅ Size is optimal"
            SETUP_ZRAM=false
        fi
    else
        echo "  ❌ Not configured"
        echo "  📊 Recommended: ${optimal_size}GB"
        needs_changes=true
        read -p "  Would you like to configure zram? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            SETUP_ZRAM=true
        else
            SETUP_ZRAM=false
        fi
    fi
    echo ""
    
    # === TMPFS CHECK ===
    echo "💾 tmpfs Cache System Analysis:"
    local tmpfs_count=$(mount | grep "tmpfs.*tmpfs-cache" | wc -l)
    if [[ $tmpfs_count -gt 0 ]]; then
        echo "  ✅ Well configured ($tmpfs_count active mounts)"
        echo "  📁 Some mounted caches:"
        mount | grep "tmpfs.*tmpfs-cache" | head -3 | awk '{print "     " $3 " (" $6 ")"}' | tr -d '()'
        [[ $tmpfs_count -gt 3 ]] && echo "     ... and $((tmpfs_count - 3)) more"
        SETUP_TMPFS=false
    else
        echo "  ❌ Not configured"
        echo "  🔍 Scanning system for cache directories..."
        
        # Analyze potential tmpfs candidates
        local cache_analysis=""
        local total_cache_size=0
        
        # Check browser caches
        for user_home in /home/*; do
            if [[ -d "$user_home" ]]; then
                local browser_cache_size=0
                for cache_dir in "$user_home/.cache/google-chrome" "$user_home/.cache/firefox" "$user_home/.cache/chromium"; do
                    if [[ -d "$cache_dir" ]]; then
                        local size=$(du -sm "$cache_dir" 2>/dev/null | cut -f1 || echo 0)
                        browser_cache_size=$((browser_cache_size + size))
                    fi
                done
                if [[ $browser_cache_size -gt 0 ]]; then
                    cache_analysis="${cache_analysis}     Browser caches: ${browser_cache_size}MB\n"
                    total_cache_size=$((total_cache_size + browser_cache_size))
                fi
            fi
        done
        
        # Check package cache
        if [[ -d /var/cache/apt ]]; then
            local apt_cache_size=$(du -sm /var/cache/apt 2>/dev/null | cut -f1 || echo 0)
            if [[ $apt_cache_size -gt 0 ]]; then
                cache_analysis="${cache_analysis}     Package cache: ${apt_cache_size}MB\n"
                total_cache_size=$((total_cache_size + apt_cache_size))
            fi
        fi
        
        # Check temp directories
        local tmp_size=$(du -sm /tmp 2>/dev/null | cut -f1 || echo 0)
        if [[ $tmp_size -gt 100 ]]; then
            cache_analysis="${cache_analysis}     Temp files: ${tmp_size}MB\n"
        fi
        
        if [[ $total_cache_size -gt 0 ]]; then
            echo "  📊 Found cacheable data:"
            echo -e "$cache_analysis"
            echo "  💡 Total cache size: ${total_cache_size}MB"
            
            # Recommend profile based on cache size and RAM
            local recommended_profile="medium"
            if [[ $ram_gb -ge 16 ]] && [[ $total_cache_size -gt 2000 ]]; then
                recommended_profile="large"
                echo "  📈 Recommended: Large profile (${ram_gb}GB RAM, ${total_cache_size}MB cache)"
            elif [[ $ram_gb -le 8 ]] || [[ $total_cache_size -lt 1000 ]]; then
                recommended_profile="small"
                echo "  📈 Recommended: Small profile (${ram_gb}GB RAM, ${total_cache_size}MB cache)"
            else
                echo "  📈 Recommended: Medium profile (${ram_gb}GB RAM, ${total_cache_size}MB cache)"
            fi
        else
            echo "  📊 No significant cache data found"
            echo "  💡 tmpfs can still improve performance for future cache usage"
        fi
        
        needs_changes=true
        read -p "  Would you like to set up tmpfs caches? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            SETUP_TMPFS=true
        else
            SETUP_TMPFS=false
        fi
    fi
    echo ""
    
    # === KERNEL PARAMS CHECK ===
    echo "⚙️ Kernel Parameters Analysis:"
    if [[ -f /etc/sysctl.d/99-system-optimization.conf ]]; then
        local swappiness=$(sysctl -n vm.swappiness 2>/dev/null)
        local dirty_ratio=$(sysctl -n vm.dirty_ratio 2>/dev/null)
        echo "  ✅ Configuration exists"
        echo "  📊 Current: swappiness=$swappiness, dirty_ratio=$dirty_ratio"
        
        # Check if optimal
        local opt_swap=5
        local opt_dirty=5
        if [[ $ram_gb -ge 16 ]]; then
            opt_swap=1; opt_dirty=3
        elif [[ $ram_gb -le 4 ]]; then
            opt_swap=10; opt_dirty=10
        fi
        
        echo "  📊 Optimal: swappiness=$opt_swap, dirty_ratio=$opt_dirty"
        
        if [[ "$swappiness" != "$opt_swap" ]] || [[ "$dirty_ratio" != "$opt_dirty" ]]; then
            echo "  ⚠️  Parameters could be optimized"
            needs_changes=true
            read -p "  Would you like to optimize kernel parameters? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                SETUP_KERNEL=true
            else
                SETUP_KERNEL=false
            fi
        else
            echo "  ✅ Parameters are optimal"
            SETUP_KERNEL=false
        fi
    else
        echo "  ❌ Not optimized"
        needs_changes=true
        read -p "  Would you like to optimize kernel parameters? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            SETUP_KERNEL=true
        else
            SETUP_KERNEL=false
        fi
    fi
    echo ""
    
    # === SERVICE CHECK ===
    echo "🔄 Systemd Service Analysis:"
    if systemctl is-enabled system-optimization.service &>/dev/null; then
        echo "  ✅ Service is installed and enabled"
        SETUP_SERVICE=false
    else
        echo "  ❌ Service not installed"
        needs_changes=true
        read -p "  Would you like to install the systemd service? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            SETUP_SERVICE=true
        else
            SETUP_SERVICE=false
        fi
    fi
    echo ""
    
    # Summary
    if [[ $needs_changes == false ]]; then
        success "🎉 All optimizations are already in great shape!"
        exit 0
    fi
    
    # Show what will be done
    local actions=()
    [[ $SETUP_ZRAM == true ]] && actions+=("zram")
    [[ $SETUP_TMPFS == true ]] && actions+=("tmpfs")
    [[ $SETUP_KERNEL == true ]] && actions+=("kernel")
    [[ $SETUP_SERVICE == true ]] && actions+=("service")
    
    if [[ ${#actions[@]} -gt 0 ]]; then
        echo "📋 Selected optimizations: ${actions[*]}"
        echo ""
        read -p "Proceed with these optimizations? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log "User cancelled optimization"
            exit 0
        fi
    else
        log "No optimizations selected"
        exit 0
    fi
    
    # Export variables for the optimization functions
    export SETUP_ZRAM SETUP_TMPFS SETUP_KERNEL SETUP_SERVICE
    export SYSTEM_RAM_GB=$ram_gb
}

setup_zram() {
    if [[ $SETUP_ZRAM != true ]]; then
        return 0
    fi
    
    log "Configuring zram..."
    
    # Calculate optimal size
    local zram_size=$((SYSTEM_RAM_GB / 2))
    [[ $zram_size -lt 2 ]] && zram_size=2
    
    # Remove existing zram properly
    if [[ -e /dev/zram0 ]]; then
        log "Removing existing zram configuration..."
        # Disable swap first
        swapoff /dev/zram0 2>/dev/null || true
        sleep 1
        
        # Reset the device
        echo 1 > /sys/block/zram0/reset 2>/dev/null || true
        sleep 1
        
        # Remove the device if possible
        if [[ -e /sys/class/zram-control/hot_remove ]]; then
            echo 0 > /sys/class/zram-control/hot_remove 2>/dev/null || true
        fi
        sleep 1
    fi
    
    # Unload and reload zram module for clean state
    rmmod zram 2>/dev/null || true
    sleep 1
    
    # Configure new zram
    log "Setting up new zram with ${zram_size}GB..."
    modprobe zram num_devices=1
    
    # Wait for device to be ready
    local attempts=0
    while [[ ! -e /dev/zram0 && $attempts -lt 10 ]]; do
        sleep 1
        ((attempts++))
    done
    
    if [[ ! -e /dev/zram0 ]]; then
        error "Failed to create zram device"
        return 1
    fi
    
    # Configure the new device
    echo "${zram_size}G" > /sys/block/zram0/disksize
    mkswap /dev/zram0
    swapon /dev/zram0
    
    success "zram reconfigured with ${zram_size}GB"
}

setup_tmpfs() {
    if [[ $SETUP_TMPFS != true ]]; then
        return 0
    fi
    
    log "Setting up tmpfs cache optimization..."
    
    # Determine profile based on RAM
    local profile="medium"
    if [[ $SYSTEM_RAM_GB -ge 16 ]]; then
        profile="large"
    elif [[ $SYSTEM_RAM_GB -le 8 ]]; then
        profile="small"
    fi
    
    # Create tmpfs directories
    mkdir -p /tmp/tmpfs-cache/{browser,ide,packages,thumbnails}
    
    # Remove existing mounts if they exist
    for dir in browser ide packages thumbnails; do
        if mount | grep -q "/tmp/tmpfs-cache/$dir"; then
            umount "/tmp/tmpfs-cache/$dir" 2>/dev/null || true
        fi
    done
    
    # Set up tmpfs mounts based on profile
    case "$profile" in
        "large")
            mount -t tmpfs -o size=4G tmpfs /tmp/tmpfs-cache/browser
            mount -t tmpfs -o size=2G tmpfs /tmp/tmpfs-cache/ide
            mount -t tmpfs -o size=3G tmpfs /tmp/tmpfs-cache/packages
            mount -t tmpfs -o size=512M tmpfs /tmp/tmpfs-cache/thumbnails
            ;;
        "medium")
            mount -t tmpfs -o size=2G tmpfs /tmp/tmpfs-cache/browser
            mount -t tmpfs -o size=1G tmpfs /tmp/tmpfs-cache/ide
            mount -t tmpfs -o size=2G tmpfs /tmp/tmpfs-cache/packages
            mount -t tmpfs -o size=256M tmpfs /tmp/tmpfs-cache/thumbnails
            ;;
        "small")
            mount -t tmpfs -o size=1G tmpfs /tmp/tmpfs-cache/browser
            mount -t tmpfs -o size=512M tmpfs /tmp/tmpfs-cache/ide
            mount -t tmpfs -o size=1G tmpfs /tmp/tmpfs-cache/packages
            mount -t tmpfs -o size=128M tmpfs /tmp/tmpfs-cache/thumbnails
            ;;
    esac
    
    # Bind mount package cache if apt is available
    if [[ -d /var/cache/apt ]]; then
        # Backup existing cache
        if [[ ! -L /var/cache/apt ]] && [[ -d /var/cache/apt ]] && [[ ! -d /var/cache/apt.backup ]]; then
            cp -a /var/cache/apt /var/cache/apt.backup
        fi
        
        # Create bind mount
        mount --bind /tmp/tmpfs-cache/packages /var/cache/apt
    fi
    
    # Add fstab entries for persistence
    local fstab_tmpfs="/etc/fstab.tmpfs"
    local browser_size ide_size packages_size thumbnails_size
    
    case "$profile" in
        "large")
            browser_size="4G"
            ide_size="2G"
            packages_size="3G"
            thumbnails_size="512M"
            ;;
        "medium")
            browser_size="2G"
            ide_size="1G"
            packages_size="2G"
            thumbnails_size="256M"
            ;;
        "small")
            browser_size="1G"
            ide_size="512M"
            packages_size="1G"
            thumbnails_size="128M"
            ;;
    esac
    
    cat > "$fstab_tmpfs" << EOF
# tmpfs cache mounts - Generated $(date)
# Profile: $profile for ${SYSTEM_RAM_GB}GB RAM
tmpfs /tmp/tmpfs-cache/browser tmpfs size=${browser_size},noatime,nosuid,nodev 0 0
tmpfs /tmp/tmpfs-cache/ide tmpfs size=${ide_size},noatime,nosuid,nodev 0 0
tmpfs /tmp/tmpfs-cache/packages tmpfs size=${packages_size},noatime,nosuid,nodev 0 0
tmpfs /tmp/tmpfs-cache/thumbnails tmpfs size=${thumbnails_size},noatime,nosuid,nodev 0 0
EOF
    
    success "tmpfs cache optimization enabled ($profile profile)"
}

tune_kernel() {
    if [[ $SETUP_KERNEL != true ]]; then
        return 0
    fi
    
    log "Optimizing kernel parameters..."
    
    local sysctl_conf="/etc/sysctl.d/99-system-optimization.conf"
    
    # Determine optimal values
    local swappiness=5
    local dirty_ratio=5
    local profile="desktop"
    
    if [[ $SYSTEM_RAM_GB -ge 16 ]]; then
        swappiness=1
        dirty_ratio=3
        profile="high-memory"
    elif [[ $SYSTEM_RAM_GB -le 4 ]]; then
        swappiness=10
        dirty_ratio=10
        profile="low-memory"
    fi
    
    cat > "$sysctl_conf" << EOF
# System Optimization - Generated $(date)
# Profile: $profile for ${SYSTEM_RAM_GB}GB RAM

# Memory management
vm.swappiness = $swappiness
vm.dirty_ratio = $dirty_ratio
vm.dirty_background_ratio = $((dirty_ratio / 2))
vm.vfs_cache_pressure = 50

# Network optimizations
net.core.rmem_max = 16777216
net.core.wmem_max = 16777216

# General performance
kernel.sched_autogroup_enabled = 1
EOF
    
    # Apply immediately
    sysctl -p "$sysctl_conf"
    success "Kernel parameters optimized ($profile profile)"
}

create_service() {
    if [[ $SETUP_SERVICE != true ]]; then
        return 0
    fi
    
    log "Creating systemd service..."
    
    # Create startup script
    cat > /usr/local/bin/system-optimization-startup.sh << 'EOF'
#!/bin/bash
# Recreate optimizations after reboot

# Ensure zram if not present
if [[ ! -e /dev/zram0 ]]; then
    ram_gb=$(free -g | awk '/^Mem:|^Speicher:/{print $2}')
    zram_size=$((ram_gb / 2))
    [[ $zram_size -lt 2 ]] && zram_size=2
    
    modprobe zram num_devices=1 2>/dev/null || true
    if [[ -e /sys/block/zram0/disksize ]]; then
        echo "${zram_size}G" > /sys/block/zram0/disksize
        mkswap /dev/zram0
        swapon /dev/zram0
    fi
fi

# Recreate tmpfs mounts if configured
if [[ -f /etc/fstab.tmpfs ]]; then
    # Create directories
    mkdir -p /tmp/tmpfs-cache/{browser,ide,packages,thumbnails}
    
    # Mount tmpfs filesystems
    while IFS= read -r line; do
        if [[ $line =~ ^tmpfs ]]; then
            # Parse the fstab line and mount
            eval "mount $line" 2>/dev/null || true
        fi
    done < /etc/fstab.tmpfs
    
    # Restore package cache bind mount
    if [[ -d /var/cache/apt ]] && [[ -d /tmp/tmpfs-cache/packages ]]; then
        mount --bind /tmp/tmpfs-cache/packages /var/cache/apt 2>/dev/null || true
    fi
fi
EOF
    
    chmod +x /usr/local/bin/system-optimization-startup.sh
    
    # Create systemd service
    cat > /etc/systemd/system/system-optimization.service << EOF
[Unit]
Description=System Performance Optimization
After=multi-user.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/system-optimization-startup.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable system-optimization.service
    
    success "Systemd service created and enabled"
}

show_final_status() {
    echo ""
    echo "🎉 Optimization Complete!"
    echo "========================"
    echo ""
    
    # Show current status
    echo "📊 Current Status:"
    
    # zram
    if [[ -e /dev/zram0 ]] && swapon --show | grep -q zram0; then
        local zram_size=$(swapon --show | grep zram0 | awk '{print $3}')
        echo "  🗜️  zram: $zram_size"
    fi
    
    # tmpfs
    local tmpfs_count=$(mount | grep "tmpfs.*tmpfs-cache" | wc -l)
    echo "  💾 tmpfs: $tmpfs_count cache mounts"
    
    # kernel
    if [[ -f /etc/sysctl.d/99-system-optimization.conf ]]; then
        local swappiness=$(sysctl -n vm.swappiness 2>/dev/null)
        echo "  ⚙️  kernel: swappiness=$swappiness"
    fi
    
    # service
    if systemctl is-enabled system-optimization.service &>/dev/null; then
        echo "  🔄 service: enabled"
    fi
    
    echo ""
    echo "🔧 Changes applied successfully!"
    echo "🔄 Reboot recommended for optimal performance"
}

main() {
    echo "🚀 Interactive System Optimizer"
    echo "==============================="
    echo ""
    
    check_root
    
    # Setup directories
    mkdir -p "$(dirname "$LOG_FILE")" "$BACKUP_DIR"
    touch "$LOG_FILE" 2>/dev/null || true
    
    # Analyze and get user input
    analyze_and_prompt
    
    # Apply selected optimizations
    setup_zram
    setup_tmpfs
    tune_kernel
    create_service
    
    # Show results
    show_final_status
    
    log "Optimization completed successfully"
}

main "$@"