#!/usr/bin/env bash

#########################################################################################
#### sysctl-tuner.sh — Apply, audit, diff, and rollback sysctl tuning profiles       ####
#### Built-in presets for web servers, database servers, and high-throughput          ####
#### workloads with automatic backup and restore                                     ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./sysctl-tuner.sh --profile web-server                                        ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
MODE=""
PROFILE="${SYSCTL_TUNER_PROFILE:-}"
CUSTOM_FILE="${SYSCTL_TUNER_CUSTOM:-}"
BACKUP_DIR="${SYSCTL_TUNER_BACKUP_DIR:-/var/lib/sysctl-tuner/backups}"
PERSIST="${SYSCTL_TUNER_PERSIST:-false}"
PERSIST_FILE="${SYSCTL_TUNER_PERSIST_FILE:-/etc/sysctl.d/99-sysctl-tuner.conf}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
RESTORE_FILE=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=$(date +%s)

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

die() {
    err "$*"
    exit 1
}

# ── Helpers ───────────────────────────────────────────────────────────
section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

require_root() {
    if [[ $EUID -ne 0 ]]; then
        die "This operation requires root privileges. Run with sudo."
    fi
}

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

# ══════════════════════════════════════════════════════════════════════
#                       BUILT-IN PROFILES
# ══════════════════════════════════════════════════════════════════════

# Each profile is a newline-separated list of "key = value" pairs.
# These are written to an associative array by load_profile().

profile_web_server() {
    cat <<'EOF'
# Web Server Profile — optimized for high connection counts and low latency
# Increase listen backlog for busy web servers
net.core.somaxconn = 4096
# TCP SYN backlog queue
net.ipv4.tcp_max_syn_backlog = 8192
# Network device backlog
net.core.netdev_max_backlog = 5000
# Reduce keepalive time (default 7200 is too long for web)
net.ipv4.tcp_keepalive_time = 600
net.ipv4.tcp_keepalive_intvl = 30
net.ipv4.tcp_keepalive_probes = 5
# Faster connection teardown
net.ipv4.tcp_fin_timeout = 15
# Reuse TIME_WAIT sockets
net.ipv4.tcp_tw_reuse = 1
# Wider ephemeral port range
net.ipv4.ip_local_port_range = 1024 65535
# SYN cookies for SYN flood protection
net.ipv4.tcp_syncookies = 1
# Increase file descriptor limit awareness
fs.file-max = 2097152
# Reduce swappiness for web workloads
vm.swappiness = 10
EOF
}

profile_db_server() {
    cat <<'EOF'
# Database Server Profile — optimized for memory-heavy, I/O-intensive workloads
# Minimize swapping — databases manage their own caches
vm.swappiness = 10
# Dirty page ratios — flush sooner for consistent write latency
vm.dirty_ratio = 15
vm.dirty_background_ratio = 5
# Dirty page expiry (centiseconds) — flush pages older than 3s
vm.dirty_expire_centisecs = 300
# Writeback interval (centiseconds)
vm.dirty_writeback_centisecs = 100
# Don't overcommit memory
vm.overcommit_memory = 0
vm.overcommit_ratio = 80
# Shared memory — allow large SHM segments for databases
kernel.shmmax = 68719476736
kernel.shmall = 4294967296
# Semaphore limits — needed by Oracle, PostgreSQL, etc.
kernel.sem = 250 32000 100 128
# Increase file descriptor limit
fs.file-max = 2097152
# Reduce TCP keepalive for connection pooling
net.ipv4.tcp_keepalive_time = 600
net.ipv4.tcp_keepalive_intvl = 30
net.ipv4.tcp_keepalive_probes = 5
# Increase listen backlog
net.core.somaxconn = 4096
EOF
}

profile_high_throughput() {
    cat <<'EOF'
# High-Throughput Profile — optimized for maximum network bandwidth
# Increase socket buffer sizes
net.core.rmem_default = 262144
net.core.wmem_default = 262144
net.core.rmem_max = 16777216
net.core.wmem_max = 16777216
# TCP buffer auto-tuning ranges (min, default, max)
net.ipv4.tcp_rmem = 4096 262144 16777216
net.ipv4.tcp_wmem = 4096 262144 16777216
# Enable TCP window scaling
net.ipv4.tcp_window_scaling = 1
# Enable BBR congestion control (requires kernel 4.9+)
net.ipv4.tcp_congestion_control = bbr
net.core.default_qdisc = fq
# Large listen backlog
net.core.somaxconn = 8192
net.ipv4.tcp_max_syn_backlog = 8192
net.core.netdev_max_backlog = 16384
# Reuse sockets and wider port range
net.ipv4.tcp_tw_reuse = 1
net.ipv4.ip_local_port_range = 1024 65535
# Faster keepalive
net.ipv4.tcp_keepalive_time = 600
net.ipv4.tcp_keepalive_intvl = 30
net.ipv4.tcp_keepalive_probes = 5
# Reduce FIN timeout
net.ipv4.tcp_fin_timeout = 10
# Enable SACK and timestamps
net.ipv4.tcp_sack = 1
net.ipv4.tcp_timestamps = 1
# Increase file descriptor limit
fs.file-max = 2097152
# Reduce swappiness
vm.swappiness = 10
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       PROFILE LOADING
# ══════════════════════════════════════════════════════════════════════

declare -a PROFILE_KEYS=()
declare -A PROFILE_VALUES=()

load_profile() {
    local source="$1"
    local raw=""

    case "$source" in
        web-server)      raw=$(profile_web_server) ;;
        db-server)       raw=$(profile_db_server) ;;
        high-throughput) raw=$(profile_high_throughput) ;;
        *)               die "Unknown built-in profile: $source (available: web-server, db-server, high-throughput)" ;;
    esac

    parse_profile_data "$raw"
}

load_custom_profile() {
    local file="$1"

    if [[ ! -f "$file" ]]; then
        die "Custom profile file not found: $file"
    fi

    local raw
    raw=$(cat "$file")
    parse_profile_data "$raw"
}

parse_profile_data() {
    local raw="$1"
    PROFILE_KEYS=()
    PROFILE_VALUES=()

    while IFS= read -r line; do
        # Skip blank lines and comments
        [[ -z "$line" ]] && continue
        [[ "$line" =~ ^[[:space:]]*# ]] && continue

        # Parse "key = value" or "key=value"
        local key value
        key=$(echo "$line" | sed 's/[[:space:]]*=.*$//' | xargs)
        value=$(echo "$line" | sed 's/^[^=]*=[[:space:]]*//' | xargs)

        if [[ -n "$key" && -n "$value" ]]; then
            PROFILE_KEYS+=("$key")
            PROFILE_VALUES["$key"]="$value"
            verbose "Loaded: $key = $value"
        fi
    done <<< "$raw"

    if [[ ${#PROFILE_KEYS[@]} -eq 0 ]]; then
        die "No valid parameters found in profile"
    fi

    verbose "Loaded ${#PROFILE_KEYS[@]} parameter(s)"
}

# ══════════════════════════════════════════════════════════════════════
#                       SYSCTL HELPERS
# ══════════════════════════════════════════════════════════════════════

get_current_value() {
    local key="$1"
    sysctl -n "$key" 2>/dev/null | xargs || echo ""
}

set_sysctl_value() {
    local key="$1"
    local value="$2"
    sysctl -w "${key}=${value}" &>/dev/null
}

# ══════════════════════════════════════════════════════════════════════
#                       BACKUP / RESTORE
# ══════════════════════════════════════════════════════════════════════

create_backup() {
    local label="${1:-manual}"

    mkdir -p "$BACKUP_DIR"

    local timestamp
    timestamp=$(date +%Y%m%d-%H%M%S)
    local backup_file="${BACKUP_DIR}/sysctl-backup-${label}-${timestamp}.conf"

    verbose "Creating backup at $backup_file"

    {
        echo "# Sysctl Tuner Backup"
        echo "# Created: $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
        echo "# Label: $label"
        echo "# Hostname: $(hostname)"
        echo ""
    } > "$backup_file"

    local count=0
    for key in "${PROFILE_KEYS[@]}"; do
        local current
        current=$(get_current_value "$key")
        if [[ -n "$current" ]]; then
            echo "${key} = ${current}" >> "$backup_file"
            ((count++)) || true
        else
            echo "# ${key} = <not set>" >> "$backup_file"
        fi
    done

    log "Backup saved: $backup_file ($count parameters)"
    echo "$backup_file"
}

get_latest_backup() {
    if [[ ! -d "$BACKUP_DIR" ]]; then
        die "No backup directory found: $BACKUP_DIR"
    fi

    local latest
    latest=$(find "$BACKUP_DIR" -maxdepth 1 -name 'sysctl-backup-*.conf' -printf '%T@\t%p\n' 2>/dev/null | sort -rn | head -1 | cut -f2)

    if [[ -z "$latest" ]]; then
        die "No backup files found in $BACKUP_DIR"
    fi

    echo "$latest"
}

do_backup() {
    load_active_profile

    require_root

    log "Backing up current sysctl values..."
    local profile_label="${PROFILE:-custom}"
    create_backup "$profile_label" > /dev/null
}

do_restore() {
    require_root

    local backup_file
    if [[ -n "$RESTORE_FILE" ]]; then
        backup_file="$RESTORE_FILE"
        if [[ ! -f "$backup_file" ]]; then
            die "Backup file not found: $backup_file"
        fi
    else
        backup_file=$(get_latest_backup)
    fi

    log "Restoring from: $backup_file"

    # Load the backup as a profile
    load_custom_profile "$backup_file"

    local applied=0
    local failed=0

    for key in "${PROFILE_KEYS[@]}"; do
        local value="${PROFILE_VALUES[$key]}"
        local current
        current=$(get_current_value "$key")

        if [[ "$current" == "$value" ]]; then
            verbose "Already set: $key = $value"
            ((applied++)) || true
            continue
        fi

        if set_sysctl_value "$key" "$value"; then
            echo -e "  ${GREEN}✓${RESET} ${key} = ${value} (was: ${current})"
            ((applied++)) || true
        else
            echo -e "  ${RED}✗${RESET} ${key} — failed to set ${value}"
            ((failed++)) || true
        fi
    done

    echo ""
    log "Restore complete: $applied applied, $failed failed ($(elapsed))"

    if [[ "$failed" -gt 0 ]]; then
        return 1
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       PROFILE RESOLUTION
# ══════════════════════════════════════════════════════════════════════

load_active_profile() {
    if [[ -n "$CUSTOM_FILE" ]]; then
        load_custom_profile "$CUSTOM_FILE"
    elif [[ -n "$PROFILE" ]]; then
        load_profile "$PROFILE"
    else
        die "No profile specified. Use --profile NAME or --custom FILE."
    fi
}

get_profile_display_name() {
    if [[ -n "$PROFILE" ]]; then
        echo "$PROFILE"
    elif [[ -n "$CUSTOM_FILE" ]]; then
        echo "custom ($(basename "$CUSTOM_FILE"))"
    else
        echo "unknown"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       MODE: DIFF
# ══════════════════════════════════════════════════════════════════════

do_diff() {
    load_active_profile

    local profile_name
    profile_name=$(get_profile_display_name)

    section_header "Diff: $profile_name"

    printf "  ${BOLD}%-38s %-14s    %-14s${RESET}\n" "PARAMETER" "CURRENT" "PROFILE"
    echo "  ─────────────────────────────────────────────────────────────"

    local would_change=0
    local already_match=0

    for key in "${PROFILE_KEYS[@]}"; do
        local target="${PROFILE_VALUES[$key]}"
        local current
        current=$(get_current_value "$key")

        if [[ -z "$current" ]]; then
            printf "  ${YELLOW}%-38s %-14s → %-14s${RESET}\n" "$key" "<not set>" "$target"
            ((would_change++)) || true
        elif [[ "$current" == "$target" ]]; then
            verbose "Match: $key = $current"
            ((already_match++)) || true
        else
            printf "  %-38s ${RED}%-14s${RESET} → ${GREEN}%-14s${RESET}\n" "$key" "$current" "$target"
            ((would_change++)) || true
        fi
    done

    echo ""
    log "$would_change parameter(s) would change, $already_match already match"
}

# ══════════════════════════════════════════════════════════════════════
#                       MODE: AUDIT
# ══════════════════════════════════════════════════════════════════════

do_audit() {
    load_active_profile

    local profile_name
    profile_name=$(get_profile_display_name)

    section_header "Audit: $profile_name"

    local pass=0
    local fail=0
    local missing=0

    for key in "${PROFILE_KEYS[@]}"; do
        local target="${PROFILE_VALUES[$key]}"
        local current
        current=$(get_current_value "$key")

        if [[ -z "$current" ]]; then
            echo -e "  ${YELLOW}?${RESET} ${key} — parameter not found"
            ((missing++)) || true
        elif [[ "$current" == "$target" ]]; then
            echo -e "  ${GREEN}✓${RESET} ${key} = ${current}"
            ((pass++)) || true
        else
            echo -e "  ${RED}✗${RESET} ${key} — expected ${target}, got ${current}"
            ((fail++)) || true
        fi
    done

    local total=$(( pass + fail + missing ))
    echo ""
    log "Audit complete: $pass passed, $fail failed, $missing missing out of $total checks ($(elapsed))"

    if [[ "$fail" -gt 0 || "$missing" -gt 0 ]]; then
        return 1
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       MODE: APPLY
# ══════════════════════════════════════════════════════════════════════

do_apply() {
    load_active_profile
    require_root

    local profile_name
    profile_name=$(get_profile_display_name)

    log "Applying profile: $profile_name"

    # Always backup before applying
    log "Backing up current values before applying..."
    local backup_file
    backup_file=$(create_backup "${PROFILE:-custom}")

    section_header "Apply: $profile_name"

    local applied=0
    local skipped=0
    local failed=0

    for key in "${PROFILE_KEYS[@]}"; do
        local target="${PROFILE_VALUES[$key]}"
        local current
        current=$(get_current_value "$key")

        if [[ "$current" == "$target" ]]; then
            verbose "Already set: $key = $target"
            ((skipped++)) || true
            continue
        fi

        if set_sysctl_value "$key" "$target"; then
            local display_current="${current:-<not set>}"
            echo -e "  ${GREEN}✓${RESET} ${key} = ${target} (was: ${display_current})"
            ((applied++)) || true
        else
            echo -e "  ${RED}✗${RESET} ${key} — failed to set ${target}"
            ((failed++)) || true
        fi
    done

    # Persist if requested
    if [[ "$PERSIST" == "true" ]]; then
        write_persist_file "$profile_name"
    fi

    echo ""
    log "Apply complete: $applied changed, $skipped unchanged, $failed failed ($(elapsed))"

    if [[ "$failed" -gt 0 ]]; then
        warn "Some parameters failed to apply. Backup saved at: $backup_file"
        return 1
    fi
}

write_persist_file() {
    local profile_name="$1"

    log "Writing persistent configuration to $PERSIST_FILE"

    {
        echo "# Sysctl Tuner — persistent configuration"
        echo "# Profile: $profile_name"
        echo "# Generated: $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
        echo "# Hostname: $(hostname)"
        echo "#"
        echo "# Applied by sysctl-tuner.sh — do not edit manually"
        echo "# To revert: sudo sysctl-tuner.sh --restore && sudo rm $PERSIST_FILE"
        echo ""
    } > "$PERSIST_FILE"

    for key in "${PROFILE_KEYS[@]}"; do
        echo "${key} = ${PROFILE_VALUES[$key]}" >> "$PERSIST_FILE"
    done

    log "Persistent config written: $PERSIST_FILE"
    log "Settings will survive reboot. Remove the file to revert."
}

# ══════════════════════════════════════════════════════════════════════
#                       MODE: EXPORT
# ══════════════════════════════════════════════════════════════════════

do_export() {
    echo "# Sysctl Export — current running values"
    echo "# Hostname: $(hostname)"
    echo "# Exported: $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
    echo "#"
    echo "# Use with: sysctl-tuner.sh --apply --custom THIS_FILE"
    echo ""

    # If a profile is specified, export only those keys
    if [[ -n "$PROFILE" || -n "$CUSTOM_FILE" ]]; then
        load_active_profile

        for key in "${PROFILE_KEYS[@]}"; do
            local current
            current=$(get_current_value "$key")
            if [[ -n "$current" ]]; then
                echo "${key} = ${current}"
            else
                echo "# ${key} = <not set>"
            fi
        done
    else
        # Export all non-default sysctl values
        sysctl -a 2>/dev/null | sort
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       BANNER
# ══════════════════════════════════════════════════════════════════════

print_banner() {
    local profile_name
    profile_name=$(get_profile_display_name)

    echo -e "${BOLD}Sysctl Tuner${RESET}"

    if [[ -n "$PROFILE" || -n "$CUSTOM_FILE" ]]; then
        echo "Profile: $profile_name"
    fi

    echo "Mode:    $MODE"
    echo "Time:    $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
}

# ══════════════════════════════════════════════════════════════════════
#                       HELP
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Modes:
  --apply              Apply a tuning profile (requires --profile or --custom)
  --audit              Compare current values against a profile
  --diff               Show what would change without applying (dry-run)
  --backup             Save current values for all keys in a profile
  --restore [FILE]     Roll back to the most recent (or specified) backup
  --export             Dump current sysctl values as a profile file

Profile options:
  --profile NAME       Built-in profile: web-server, db-server, high-throughput
  --custom FILE        Load a custom profile file (key=value format)

Options:
  --persist            Write drop-in to /etc/sysctl.d/ for reboot persistence
  --backup-dir DIR     Backup directory (default: /var/lib/sysctl-tuner/backups)
  --verbose            Debug output
  --no-color           Disable colored output
  --help               Show this help message

Environment variables:
  SYSCTL_TUNER_PROFILE       Built-in profile name
  SYSCTL_TUNER_CUSTOM        Path to custom profile file
  SYSCTL_TUNER_BACKUP_DIR    Backup directory
  SYSCTL_TUNER_PERSIST       Set to "true" to persist
  SYSCTL_TUNER_PERSIST_FILE  Drop-in file path
  VERBOSE                    Set to "true" for debug output
  COLOR                      auto, always, never

Examples:
  # Preview changes for web-server profile
  sudo $SCRIPT_NAME --diff --profile web-server

  # Apply database profile with persistence
  sudo $SCRIPT_NAME --apply --profile db-server --persist

  # Audit current system against high-throughput profile
  $SCRIPT_NAME --audit --profile high-throughput

  # Export current values as a profile
  $SCRIPT_NAME --export > my-profile.conf

  # Apply a custom profile
  sudo $SCRIPT_NAME --apply --custom my-profile.conf

  # Roll back the last change
  sudo $SCRIPT_NAME --restore
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --apply)
                MODE="apply"
                shift
                ;;
            --audit)
                MODE="audit"
                shift
                ;;
            --diff)
                MODE="diff"
                shift
                ;;
            --backup)
                MODE="backup"
                shift
                ;;
            --restore)
                MODE="restore"
                shift
                # Optional: next arg could be a backup file path
                if [[ $# -gt 0 && ! "$1" =~ ^-- ]]; then
                    RESTORE_FILE="$1"
                    shift
                fi
                ;;
            --export)
                MODE="export"
                shift
                ;;
            --profile)
                shift
                if [[ $# -eq 0 ]]; then
                    die "--profile requires a name (web-server, db-server, high-throughput)"
                fi
                PROFILE="$1"
                shift
                ;;
            --custom)
                shift
                if [[ $# -eq 0 ]]; then
                    die "--custom requires a file path"
                fi
                CUSTOM_FILE="$1"
                shift
                ;;
            --persist)
                PERSIST="true"
                shift
                ;;
            --backup-dir)
                shift
                if [[ $# -eq 0 ]]; then
                    die "--backup-dir requires a directory path"
                fi
                BACKUP_DIR="$1"
                shift
                ;;
            --verbose)
                VERBOSE="true"
                shift
                ;;
            --no-color)
                COLOR="never"
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                die "Unknown option: $1 (see --help)"
                ;;
        esac
    done

    if [[ -z "$MODE" ]]; then
        # Default to diff if a profile is given but no mode
        if [[ -n "$PROFILE" || -n "$CUSTOM_FILE" ]]; then
            MODE="diff"
        else
            show_help
            exit 0
        fi
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    setup_colors

    # Export mode goes straight to stdout — no banner
    if [[ "$MODE" == "export" ]]; then
        do_export
        exit 0
    fi

    print_banner

    case "$MODE" in
        apply)   do_apply   ;;
        audit)   do_audit   ;;
        diff)    do_diff    ;;
        backup)  do_backup  ;;
        restore) do_restore ;;
        *)       die "Unknown mode: $MODE" ;;
    esac
}

main "$@"