#!/usr/bin/zsh

# Define Named Terminal Colors
BLUE=$'\e[34m'
GREEN=$'\e[32m'
BOLD=$'\e[1m'
RESET=$'\e[0m'

# Formatting function for RAM/Swap (uses awk for float math)
format_size() {
    local mb_val=$1
    if [ "$mb_val" -ge 4096 ]; then
        echo "$(echo "$mb_val" | awk '{printf "%.2f", $1 / 1024}') GiB"
    else
        echo "${mb_val} MiB"
    fi
}

# Gather Data
OS=$(grep '^PRETTY_NAME' /etc/os-release | cut -d'=' -f2 | tr -d '"')
KERNEL=$(uname -r)
DEVICE=$(cat /sys/devices/virtual/dmi/id/product_name 2>/dev/null || hostname)
CPU=$(grep "model name" /proc/cpuinfo | head -1 | cut -d ":" -f 2 | sed 's/^[ \t]*//')
GPU=$(lspci | grep -i 'vga\|3d\|2d' | cut -d ":" -f 3 | sed 's/^[ \t]*//' | head -n 1)
SHELL_NAME=$(basename "$SHELL")

# Memory & Swap (Raw MB)
MEM_USED_MB=$(free -m | awk '/Mem:/ {print $3}')
MEM_TOTAL_MB=$(free -m | awk '/Mem:/ {print $2}')
SWAP_USED_MB=$(free -m | awk '/Swap:/ {print $3}')
SWAP_TOTAL_MB=$(free -m | awk '/Swap:/ {print $2}')

# Format for output
MEM_STR="$(format_size $MEM_USED_MB) / $(format_size $MEM_TOTAL_MB)"
SWAP_STR="$(format_size $SWAP_USED_MB) / $(format_size $SWAP_TOTAL_MB)"

# Network
LOCAL_IP=$(hostname -I | awk '{print $1}')

# Get disk info (filter out loop and 0B disks)
DISK_DATA=$(lsblk -dno NAME,SIZE | grep -v "loop" | awk '$2 != "0B"')
DISK_COUNT=$(echo "$DISK_DATA" | grep -c .)

# Output
echo ""
echo "${BOLD}>> System Status for ${USER}${RESET}"
echo "${BOLD}--------------------------------------------------${RESET}"

printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "OS:" "$OS"
printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "Device:" "$DEVICE"
printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "Kernel:" "$KERNEL"
printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "Shell:" "$SHELL_NAME"
printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "CPU:" "$CPU"
printf "  ${BOLD}${BLUE}%-15s${RESET} %s\n" "GPU:" "${GPU:-Unknown}"

echo "${BOLD}--------------------------------------------------${RESET}"

printf "  ${BOLD}${GREEN}%-15s${RESET} %s\n" "Memory:" "$MEM_STR"
printf "  ${BOLD}${GREEN}%-15s${RESET} %s\n" "Swap:" "$SWAP_STR"
printf "  ${BOLD}${GREEN}%-15s${RESET} %s\n" "Local IP:" "$LOCAL_IP"

# Disk output: single line if 1 disk, list if multiple
if [ "$DISK_COUNT" -eq 1 ]; then
    read -r name size <<< "$DISK_DATA"
    USAGE=$(df -h -x squashfs -x fuse -x tmpfs 2>/dev/null | grep "/dev/$name" | awk '{print $5}' | head -n 1)
    [ -z "$USAGE" ] && USAGE="N/A"
    printf "  ${BOLD}${GREEN}%-15s${RESET} %s %s (%s)\n" "Disk:" "$name" "$size" "$USAGE"
else
    echo "  ${BOLD}${GREEN}Disks:${RESET}"
    echo "$DISK_DATA" | while read -r name size; do
        USAGE=$(df -h -x squashfs -x fuse -x tmpfs 2>/dev/null | grep "/dev/$name" | awk '{print $5}' | head -n 1)
        [ -z "$USAGE" ] && USAGE="N/A"
        printf "    [-] %-10s Size: %-8s Usage: %s\n" "$name" "$size" "$USAGE"
    done
fi

echo ""