#!/bin/bash

# ----- Colors -----
RED="\e[31m"
GREEN="\e[32m"
YELLOW="\e[33m"
BLUE="\e[34m"
MAGENTA="\e[35m"
CYAN="\e[36m"
WHITE="\e[97m"
RESET="\e[0m"

# ----- Location (quick, won't hang if offline) -----
CITY=$(curl -sS -m 1 https://ipinfo.io/city || true)
COUNTRY=$(curl -sS -m 1 https://ipinfo.io/country || true)
CITY=${CITY:-Unknown}
COUNTRY=${COUNTRY:-Unknown}

# ----- Workload Tag Detection -----
# Prefer cached value if you later add a systemd timer that writes /run/host-tag
if [ -r /run/host-tag ]; then
  TAG=$(cat /run/host-tag)
else
  # Detect by processes (no false positive from grep itself)
  if pgrep -fa 'runpod|com\.runpod' >/dev/null 2>&1; then
    TAG="Runpod"
  elif pgrep -fa 'vastai|vast\.ai' >/dev/null 2>&1; then
    TAG="Vastai"
  elif pgrep -fa 'gpu-burn|gpu_burn' >/dev/null 2>&1; then
    TAG="Burn"
  else
    TAG="Not Assigned"
  fi
fi

# Color for tag
case "$TAG" in
  runpod) TAG_COLOR="$CYAN" ;;
  vastai) TAG_COLOR="$MAGENTA" ;;
  burn)   TAG_COLOR="$YELLOW" ;;
  *)      TAG_COLOR="$BLUE" ;; # idle/other
esac

# ----- MOTD -----
echo -e "------------------------------------------"
echo -e "  Welcome - ${GREEN}$(hostname)${RESET}  [${TAG_COLOR}${TAG}${RESET}]"
echo -e "  Location: ${COUNTRY}${RESET}"
echo -e "  Server Time: ${GREEN}$(date)${RESET}"
echo -e "  Logged in as: ${GREEN}$(whoami)${RESET}"
echo -e "------------------------------------------"