#!/bin/bash
# Linux System Tuning Suite - Main Orchestrator
# Intelligent system optimization based on hardware and usage analysis

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODULES_DIR="$SCRIPT_DIR/modules"
PROFILES_DIR="$SCRIPT_DIR/profiles"
CONFIG_DIR="$SCRIPT_DIR/configs"
BACKUP_DIR="/var/lib/system-tuning/backups"

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

# Global variables
DRY_RUN=false
VERBOSE=false
AUTO_MODE=false
PROFILE=""
FORCE=false

log() {
    echo -e "${BLUE}[$(date '+%H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

show_usage() {
    cat << EOF
Linux System Tuning Suite

Usage: $0 [OPTIONS] [ACTIONS]

OPTIONS:
    -h, --help          Show this help message
    -v, --verbose       Enable verbose output
    -n, --dry-run       Show what would be done without making changes
    -f, --force         Force operations without confirmation
    -p, --profile NAME  Use specific optimization profile

ACTIONS:
    --analyze           Run system analysis only
    --auto              Automatic optimization based on system analysis
    --enable-zram       Configure zram compression
    --enable-tmpfs      Set up tmpfs optimizations
    --tune-kernel       Optimize kernel parameters
    --create-service    Install systemd service for persistent optimizations
    --benchmark         Run performance benchmarks
    --rollback          Restore previous configuration
    --status            Show current optimization status

PROFILES:
    desktop             Desktop/workstation optimization (default)
    server              Server optimization
    gaming              Gaming-focused optimization
    development         Development workstation optimization
    minimal             Minimal safe optimizations

EXAMPLES:
    $0 --analyze                    # Analyze system only
    $0 --auto                       # Automatic optimization
    $0 --profile gaming --auto      # Gaming-optimized setup
    $0 --dry-run --auto             # Show what would be done
    $0 --enable-zram --enable-tmpfs # Specific optimizations

EOF
}

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

setup_directories() {
    mkdir -p "$BACKUP_DIR"
    mkdir -p /var/lib/system-tuning/configs
    mkdir -p /var/lib/system-tuning/logs
}

load_modules() {
    for module in "$MODULES_DIR"/*.sh; do
        if [[ -f "$module" ]]; then
            source "$module"
            [[ "$VERBOSE" == "true" ]] && log "Loaded module: $(basename "$module")"
        fi
    done
}

create_backup() {
    local backup_name="backup-$(date +%Y%m%d-%H%M%S)"
    local backup_path="$BACKUP_DIR/$backup_name"
    
    mkdir -p "$backup_path"
    
    log "Creating system backup..."
    
    # Backup important configs
    [[ -f /etc/sysctl.conf ]] && cp /etc/sysctl.conf "$backup_path/"
    [[ -f /etc/fstab ]] && cp /etc/fstab "$backup_path/"
    
    # Save current state
    free -h > "$backup_path/memory-before.txt"
    mount > "$backup_path/mounts-before.txt"
    sysctl -a > "$backup_path/sysctl-before.txt" 2>/dev/null
    
    echo "$backup_name" > /var/lib/system-tuning/last-backup
    success "Backup created: $backup_path"
}

analyze_system() {
    log "Running system analysis..."
    
    if [[ -f "$SCRIPT_DIR/system-analyzer.sh" ]]; then
        bash "$SCRIPT_DIR/system-analyzer.sh"
    else
        error "System analyzer not found"
        return 1
    fi
}

auto_optimize() {
    log "Starting automatic optimization..."
    
    # Run analysis first
    analyze_system
    
    # Determine optimal configuration based on hardware
    local ram_gb=$(free -g | awk '/^Mem:/{print $2}')
    local cpu_cores=$(nproc)
    
    log "Detected: ${ram_gb}GB RAM, ${cpu_cores} CPU cores"
    
    # Apply optimizations based on available RAM
    if [[ $ram_gb -ge 16 ]]; then
        log "High-memory system detected (≥16GB) - applying aggressive optimizations"
        enable_zram "12G"
        enable_tmpfs_cache "large"
        tune_kernel_params "desktop-high-memory"
    elif [[ $ram_gb -ge 8 ]]; then
        log "Medium-memory system detected (8-15GB) - applying moderate optimizations"
        enable_zram "6G"
        enable_tmpfs_cache "medium"
        tune_kernel_params "desktop"
    elif [[ $ram_gb -ge 4 ]]; then
        log "Low-memory system detected (4-7GB) - applying conservative optimizations"
        enable_zram "4G"
        enable_tmpfs_cache "small"
        tune_kernel_params "desktop-low-memory"
    else
        warn "Very low memory system (<4GB) - applying minimal optimizations only"
        enable_zram "2G"
        tune_kernel_params "minimal"
    fi
    
    create_systemd_service
    success "Automatic optimization complete!"
}

enable_zram() {
    local size=${1:-"8G"}
    log "Configuring zram with size: $size"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        echo "[DRY RUN] Would configure zram with size $size"
        return
    fi
    
    # Check if zram module exists
    if ! modinfo zram >/dev/null 2>&1; then
        error "zram module not available"
        return 1
    fi
    
    # Configure zram
    if [[ -e /dev/zram0 ]]; then
        swapoff /dev/zram0 2>/dev/null || true
        echo 1 > /sys/class/zram-control/hot_remove 2>/dev/null || true
    fi
    
    modprobe zram num_devices=1
    echo "$size" > /sys/block/zram0/disksize
    mkswap /dev/zram0
    swapon /dev/zram0
    
    success "zram configured with $size"
}

enable_tmpfs_cache() {
    local profile=${1:-"medium"}
    log "Setting up tmpfs cache optimization (profile: $profile)"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        echo "[DRY RUN] Would set up tmpfs cache with profile: $profile"
        return
    fi
    
    # Create tmpfs directories
    mkdir -p /tmp/tmpfs-cache/{browser,ide,packages,thumbnails}
    
    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 [[ -d /var/cache/apt ]]; then
        mount --bind /tmp/tmpfs-cache/packages /var/cache/apt
    fi
    
    success "tmpfs cache optimization enabled"
}

tune_kernel_params() {
    local profile=${1:-"desktop"}
    log "Tuning kernel parameters (profile: $profile)"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        echo "[DRY RUN] Would tune kernel parameters with profile: $profile"
        return
    fi
    
    local sysctl_conf="/etc/sysctl.d/99-system-tuning.conf"
    
    cat > "$sysctl_conf" << EOF
# System Tuning - Generated $(date)
# Profile: $profile

EOF
    
    case "$profile" in
        "desktop-high-memory")
            cat >> "$sysctl_conf" << EOF
# High-memory desktop optimization
vm.swappiness = 1
vm.dirty_ratio = 3
vm.dirty_background_ratio = 1
vm.vfs_cache_pressure = 40
vm.page-cluster = 0
EOF
            ;;
        "desktop")
            cat >> "$sysctl_conf" << EOF
# Desktop optimization
vm.swappiness = 5
vm.dirty_ratio = 5
vm.dirty_background_ratio = 2
vm.vfs_cache_pressure = 50
vm.page-cluster = 0
EOF
            ;;
        "desktop-low-memory")
            cat >> "$sysctl_conf" << EOF
# Low-memory desktop optimization
vm.swappiness = 10
vm.dirty_ratio = 10
vm.dirty_background_ratio = 5
vm.vfs_cache_pressure = 100
EOF
            ;;
        "minimal")
            cat >> "$sysctl_conf" << EOF
# Minimal safe optimization
vm.swappiness = 10
vm.dirty_ratio = 10
vm.dirty_background_ratio = 5
EOF
            ;;
    esac
    
    # Apply immediately
    sysctl -p "$sysctl_conf"
    success "Kernel parameters tuned"
}

create_systemd_service() {
    log "Creating systemd service for persistent optimizations"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        echo "[DRY RUN] Would create systemd service"
        return
    fi
    
    # Create startup script
    cat > /usr/local/bin/system-tuning-startup.sh << 'EOF'
#!/bin/bash
# System Tuning Startup Script
# Auto-generated by Linux System Tuning Suite

log() {
    echo "[$(date '+%H:%M:%S')] $1" | systemd-cat -t system-tuning
}

log "Starting system optimizations..."

# Recreate tmpfs mounts
if [[ -d /tmp/tmpfs-cache ]]; then
    for cache_dir in /tmp/tmpfs-cache/*; do
        if [[ -d "$cache_dir" ]] && ! mountpoint -q "$cache_dir"; then
            size=$(df -h "$cache_dir" 2>/dev/null | awk 'NR==2{print $2}' || echo "1G")
            mount -t tmpfs -o size="$size" tmpfs "$cache_dir"
            log "Mounted tmpfs: $cache_dir ($size)"
        fi
    done
fi

# Ensure zram is configured
if [[ ! -e /dev/zram0 ]]; then
    modprobe zram num_devices=1 2>/dev/null || true
    if [[ -e /sys/block/zram0/disksize ]]; then
        echo "8G" > /sys/block/zram0/disksize
        mkswap /dev/zram0
        swapon /dev/zram0
        log "Configured zram swap"
    fi
fi

log "System optimizations applied"
EOF

    chmod +x /usr/local/bin/system-tuning-startup.sh
    
    # Create systemd service
    cat > /etc/systemd/system/system-tuning.service << EOF
[Unit]
Description=System Performance Tuning
After=multi-user.target
Wants=multi-user.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/system-tuning-startup.sh
RemainAfterExit=yes
StandardOutput=journal
StandardError=journal

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

show_status() {
    echo "=== SYSTEM TUNING STATUS ==="
    echo
    
    echo "💾 Memory:"
    free -h
    echo
    
    echo "🔄 Swap/zram:"
    if command -v zramctl >/dev/null 2>&1; then
        zramctl 2>/dev/null || echo "No zram devices"
    else
        echo "zramctl not available"
    fi
    echo
    
    echo "📁 tmpfs mounts:"
    df -h | grep tmpfs | head -10
    echo
    
    echo "⚙️ Key kernel parameters:"
    echo "  Swappiness: $(cat /proc/sys/vm/swappiness)"
    echo "  Dirty ratio: $(cat /proc/sys/vm/dirty_ratio)" 
    echo "  VFS cache pressure: $(cat /proc/sys/vm/vfs_cache_pressure)"
    echo
    
    echo "🔧 Service status:"
    if systemctl is-enabled system-tuning.service >/dev/null 2>&1; then
        echo "  System tuning service: ENABLED"
    else
        echo "  System tuning service: DISABLED"
    fi
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_usage
                exit 0
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -n|--dry-run)
                DRY_RUN=true
                shift
                ;;
            -f|--force)
                FORCE=true
                shift
                ;;
            -p|--profile)
                PROFILE="$2"
                shift 2
                ;;
            --analyze)
                analyze_system
                exit 0
                ;;
            --auto)
                AUTO_MODE=true
                shift
                ;;
            --enable-zram)
                enable_zram
                shift
                ;;
            --enable-tmpfs)
                enable_tmpfs_cache
                shift
                ;;
            --tune-kernel)
                tune_kernel_params
                shift
                ;;
            --create-service)
                create_systemd_service
                shift
                ;;
            --status)
                show_status
                exit 0
                ;;
            *)
                error "Unknown option: $1"
                show_usage
                exit 1
                ;;
        esac
    done
}

main() {
    echo "🚀 Linux System Tuning Suite"
    echo "=============================="
    
    check_root
    setup_directories
    
    # Load modules if available
    if [[ -d "$MODULES_DIR" ]]; then
        load_modules
    fi
    
    # Parse arguments
    parse_args "$@"
    
    # If no specific action, show usage
    if [[ "$AUTO_MODE" != "true" ]] && [[ $# -eq 0 ]]; then
        show_usage
        exit 0
    fi
    
    # Create backup before making changes
    if [[ "$DRY_RUN" != "true" ]]; then
        create_backup
    fi
    
    # Execute auto mode if requested
    if [[ "$AUTO_MODE" == "true" ]]; then
        auto_optimize
    fi
    
    success "Operation completed successfully!"
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi