#!/usr/bin/env bash

#########################################################################################
#### motd-generator.sh — Generate a dynamic MOTD with system stats and health info   ####
#### Shows hostname, IP, uptime, disk, load, memory, updates, and service status     ####
#### Dry-run by default for --install — use --force to write to update-motd.d        ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.02                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./motd-generator.sh                                                           ####
####   ./motd-generator.sh --plain                                                   ####
####   ./motd-generator.sh --install --force                                         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
DRY_RUN="${DRY_RUN:-true}"
MODE="${MODE:-display}"
PLAIN="${PLAIN:-false}"
MOTD_TARGET="${MOTD_TARGET:-/etc/update-motd.d/99-custom}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME

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

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${CYAN}[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; }

# ── Data Gathering ───────────────────────────────────────────────────

get_hostname() {
    hostname -f 2>/dev/null || hostname
}

get_primary_ip() {
    local ip=""
    if command -v ip &>/dev/null; then
        ip=$(ip route get 1.1.1.1 2>/dev/null | grep -oP 'src \K[\d.]+' | head -1)
    fi
    if [[ -z "$ip" ]] && command -v hostname &>/dev/null; then
        ip=$(hostname -I 2>/dev/null | awk '{print $1}')
    fi
    echo "${ip:-N/A}"
}

get_uptime() {
    uptime -p 2>/dev/null || uptime | sed 's/.*up //' | sed 's/,  [0-9]* user.*//'
}

get_disk_usage() {
    df / 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%'
}

get_load_average() {
    cut -d' ' -f1-3 /proc/loadavg 2>/dev/null || echo "N/A"
}

get_memory_usage() {
    if [[ -f /proc/meminfo ]]; then
        local total_kb avail_kb
        total_kb=$(awk '/^MemTotal:/ {print $2}' /proc/meminfo)
        avail_kb=$(awk '/^MemAvailable:/ {print $2}' /proc/meminfo)
        if [[ "$total_kb" -gt 0 ]]; then
            local pct
            pct=$(awk "BEGIN { printf \"%.0f\", ($total_kb - $avail_kb) * 100 / $total_kb }")
            local total_mb=$(( total_kb / 1024 ))
            local used_mb=$(( (total_kb - avail_kb) / 1024 ))
            echo "${pct}|${used_mb}|${total_mb}"
            return
        fi
    fi
    echo "0|0|0"
}

get_pending_updates() {
    if command -v apt-get &>/dev/null; then
        apt list --upgradable 2>/dev/null | grep -c "upgradable" || true
    elif command -v dnf &>/dev/null; then
        dnf check-update --quiet 2>/dev/null | grep -cE "^\S" || true
    elif command -v yum &>/dev/null; then
        yum check-update --quiet 2>/dev/null | grep -cE "^\S" || true
    else
        echo "N/A"
    fi
}

get_logged_in_users() {
    who 2>/dev/null | wc -l
}

get_failed_services() {
    if command -v systemctl &>/dev/null; then
        systemctl --no-legend --state=failed 2>/dev/null | wc -l
    else
        echo "N/A"
    fi
}

# ── Color Thresholds ─────────────────────────────────────────────────

threshold_color() {
    local value="$1"
    local warn_at="${2:-75}"
    local crit_at="${3:-90}"

    if [[ "$value" -ge "$crit_at" ]]; then
        echo "$RED"
    elif [[ "$value" -ge "$warn_at" ]]; then
        echo "$YELLOW"
    else
        echo "$GREEN"
    fi
}

load_color() {
    local load_1m="$1"
    local cpus
    cpus=$(nproc 2>/dev/null || echo 1)
    local pct
    pct=$(awk "BEGIN { printf \"%.0f\", ($load_1m / $cpus) * 100 }")
    threshold_color "$pct" 75 90
}

# ── MOTD Output ──────────────────────────────────────────────────────

generate_motd_plain() {
    local host ip up disk_pct load mem_info mem_pct mem_used mem_total
    local updates users_count failed_count

    host=$(get_hostname)
    ip=$(get_primary_ip)
    up=$(get_uptime)
    disk_pct=$(get_disk_usage)
    load=$(get_load_average)
    mem_info=$(get_memory_usage)
    mem_pct=$(echo "$mem_info" | cut -d'|' -f1)
    mem_used=$(echo "$mem_info" | cut -d'|' -f2)
    mem_total=$(echo "$mem_info" | cut -d'|' -f3)
    updates=$(get_pending_updates)
    users_count=$(get_logged_in_users)
    failed_count=$(get_failed_services)

    local dc mc lc
    dc=$(threshold_color "$disk_pct" 75 90)
    mc=$(threshold_color "$mem_pct" 75 90)
    local load_1m
    load_1m=$(echo "$load" | awk '{print $1}')
    lc=$(load_color "$load_1m")

    echo -e "  ${BOLD}Hostname:${RESET}          $host"
    echo -e "  ${BOLD}IP Address:${RESET}        $ip"
    echo -e "  ${BOLD}Uptime:${RESET}            $up"
    echo -e "  ${BOLD}Disk (root):${RESET}       ${dc}${disk_pct}%${RESET}"
    echo -e "  ${BOLD}Load Average:${RESET}      ${lc}${load}${RESET}"
    echo -e "  ${BOLD}Memory:${RESET}            ${mc}${mem_used}M / ${mem_total}M (${mem_pct}%)${RESET}"
    echo -e "  ${BOLD}Pending Updates:${RESET}   $updates"
    echo -e "  ${BOLD}Logged-in Users:${RESET}   $users_count"
    echo -e "  ${BOLD}Failed Services:${RESET}   $failed_count"
}

generate_motd_box() {
    local host ip up disk_pct load mem_info mem_pct mem_used mem_total
    local updates users_count failed_count

    host=$(get_hostname)
    ip=$(get_primary_ip)
    up=$(get_uptime)
    disk_pct=$(get_disk_usage)
    load=$(get_load_average)
    mem_info=$(get_memory_usage)
    mem_pct=$(echo "$mem_info" | cut -d'|' -f1)
    mem_used=$(echo "$mem_info" | cut -d'|' -f2)
    mem_total=$(echo "$mem_info" | cut -d'|' -f3)
    updates=$(get_pending_updates)
    users_count=$(get_logged_in_users)
    failed_count=$(get_failed_services)

    local dc mc lc
    dc=$(threshold_color "$disk_pct" 75 90)
    mc=$(threshold_color "$mem_pct" 75 90)
    local load_1m
    load_1m=$(echo "$load" | awk '{print $1}')
    lc=$(load_color "$load_1m")

    local fc_color="$GREEN"
    if [[ "$failed_count" != "N/A" && "$failed_count" -gt 0 ]]; then
        fc_color="$RED"
    fi

    # Build rows as "label|value" pairs to measure widest content
    local label_w=18
    local rows=(
        "Hostname:|$host"
        "IP Address:|$ip"
        "Uptime:|$up"
        "Disk (root):|${disk_pct}%"
        "Load Average:|$load"
        "Memory:|${mem_used}M / ${mem_total}M (${mem_pct}%)"
        "Pending Updates:|$updates"
        "Logged-in Users:|$users_count"
        "Failed Services:|$failed_count"
    )

    # Calculate box width from widest content
    local header="System Status: ${host}"
    local dateline
    dateline=$(date '+%Y-%m-%d %H:%M:%S %Z')
    local w=${#header}
    [[ ${#dateline} -gt $w ]] && w=${#dateline}

    local row label value row_len
    for row in "${rows[@]}"; do
        label="${row%%|*}"
        value="${row#*|}"
        row_len=$(( label_w + 1 + ${#value} ))
        [[ $row_len -gt $w ]] && w=$row_len
    done

    # Add 2 for inner padding (space on each side)
    w=$(( w + 2 ))
    # Minimum width
    [[ $w -lt 56 ]] && w=56

    # Build box-drawing borders at calculated width
    local bar
    bar=$(printf '═%.0s' $(seq 1 $((w + 2))))
    local border="╔${bar}╗"
    local bottom="╚${bar}╝"
    local sep="╠${bar}╣"

    BOX_W=$w  # export to _box_row

    echo ""
    echo -e "  ${CYAN}${border}${RESET}"
    printf "  ${CYAN}║${RESET} ${BOLD}%-${w}s${RESET} ${CYAN}║${RESET}\n" "$header"
    printf "  ${CYAN}║${RESET} ${DIM}%-${w}s${RESET} ${CYAN}║${RESET}\n" "$dateline"
    echo -e "  ${CYAN}${sep}${RESET}"
    _box_row "Hostname:         " "$host"
    _box_row "IP Address:       " "$ip"
    _box_row "Uptime:           " "$up"
    _box_row "Disk (root):      " "${disk_pct}%" "$dc"
    _box_row "Load Average:     " "$load" "$lc"
    _box_row "Memory:           " "${mem_used}M / ${mem_total}M (${mem_pct}%)" "$mc"
    _box_row "Pending Updates:  " "$updates"
    _box_row "Logged-in Users:  " "$users_count"
    _box_row "Failed Services:  " "$failed_count" "$fc_color"
    echo -e "  ${CYAN}${bottom}${RESET}"
    echo ""
}

_box_row() {
    # Print a row inside the box with correct padding regardless of content length
    # Usage: _box_row "Label:" "value" [color_prefix]
    local label="$1"
    local value="$2"
    local color="${3:-}"
    local inner_w="${BOX_W:-56}"

    # Build the visible text (no ANSI)
    local vis_text=" ${label} ${value}"
    local vis_len=${#vis_text}

    # Pad to fill the box
    local pad_len=$(( inner_w - vis_len ))
    [[ $pad_len -lt 0 ]] && pad_len=0
    local padding
    padding=$(printf '%*s' "$pad_len" '')

    # Build output with optional color on value
    if [[ -n "$color" ]]; then
        printf "  ${CYAN}║${RESET}  %s %b%s${RESET}%s ${CYAN}║${RESET}\n" "$label" "$color" "$value" "$padding"
    else
        printf "  ${CYAN}║${RESET}  %s %s%s ${CYAN}║${RESET}\n" "$label" "$value" "$padding"
    fi
}

# ── Install ──────────────────────────────────────────────────────────

install_motd() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would install MOTD script to ${MOTD_TARGET}"
        log "[DRY-RUN] Run with --force to actually install"
        echo ""
        log "Generated script preview:"
        echo "  #!/bin/bash"
        echo "  # Generated by ${SCRIPT_NAME} on $(date '+%Y-%m-%d %H:%M:%S')"
        echo "  $(readlink -f "$0") --plain --no-color"
        return
    fi

    if [[ $EUID -ne 0 ]]; then
        err "Installation requires root privileges"
        exit 1
    fi

    local script_path
    script_path=$(readlink -f "$0")

    local motd_dir
    motd_dir=$(dirname "$MOTD_TARGET")
    if [[ ! -d "$motd_dir" ]]; then
        mkdir -p "$motd_dir"
    fi

    cat > "$MOTD_TARGET" <<MOTD_EOF
#!/bin/bash
# Generated by ${SCRIPT_NAME} on $(date '+%Y-%m-%d %H:%M:%S')
# Source: ${script_path}
${script_path} --plain
MOTD_EOF

    chmod 755 "$MOTD_TARGET"
    log "Installed MOTD script to ${MOTD_TARGET}"
    log "The MOTD will update on each login"
}

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Generate a dynamic MOTD with system stats and health info

USAGE:
  ${SCRIPT_NAME} [OPTIONS]

OPTIONS:
  --plain               Plain text output (no box art)
  --install             Install as ${MOTD_TARGET} (dry-run by default)
  --force               Actually install (disables dry-run)
  --verbose             Enable debug output
  --no-color            Disable colored output
  --help                Show this help

ENVIRONMENT VARIABLES:
  COLOR                 Color mode: auto, always, never (default: auto)
  MOTD_TARGET           Install target (default: ${MOTD_TARGET})

EXAMPLES:
  # Display MOTD with box art
  ./motd-generator.sh

  # Plain format for piping
  ./motd-generator.sh --plain --no-color

  # Preview install (dry-run)
  ./motd-generator.sh --install

  # Actually install
  sudo ./motd-generator.sh --install --force
EOF
}

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

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --plain)
                PLAIN="true"; shift ;;
            --install)
                MODE="install"; shift ;;
            --force)
                DRY_RUN="false"; shift ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                setup_colors
                usage
                exit 0 ;;
            *)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 1 ;;
        esac
    done
}

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

main() {
    parse_args "$@"
    setup_colors

    case "$MODE" in
        install)
            install_motd
            ;;
        display)
            if [[ "$PLAIN" == "true" ]]; then
                generate_motd_plain
            else
                generate_motd_box
            fi
            ;;
    esac
}

main "$@"