#!/usr/bin/env bash
# Control of the TensorCash worker (Aria Energy 1.7)
cd "$(dirname "$0")" || exit 1
ENERGY_VERSION="1.7"
aria_logo() {   # Aria Energy logo: ARIA + ENERGY in white, orange lightning bolt (website palette); $1 = version
  local O W WB GR Z
  if [ -t 1 ] && [ -z "${NO_COLOR:-}" ]; then O=$'\e[1;38;2;255;77;0m'; W=$'\e[38;2;237;232;219m'; WB=$'\e[1;38;2;237;232;219m'; GR=$'\e[38;2;163;157;144m'; Z=$'\e[0m'; else O=""; W=""; WB=""; GR=""; Z=""; fi
  echo
  echo "${WB}   █████╗ ██████╗ ██╗ █████╗ ${Z}       ${O}▄▀${Z}"
  echo "${WB}  ██╔══██╗██╔══██╗██║██╔══██╗${Z}     ${O}▄▀${Z}       ${WB}Aria Energy${Z}"
  echo "${WB}  ███████║██████╔╝██║███████║${Z}   ${O}▄▀▀▀▀▀${Z}     ${W}TensorCash · Proof v4 · ARIA engine${Z}"
  echo "${WB}  ██╔══██║██╔══██╗██║██╔══██║${Z}      ${O}▄▀${Z}      ${GR}${1:-1.3}${Z}"
  echo "${WB}  ██║  ██║██║  ██║██║██║  ██║${Z}    ${O}▄▀${Z}"
  echo "${WB}  ╚═╝  ╚═╝╚═╝  ╚═╝╚═╝╚═╝  ╚═╝${Z}   ${O}▀${Z}"
  echo "${WB}         E N E R G Y${Z}"
  echo
}
aria_fiche() {   # fields under the logo (address shortened 10…8, 1% fee in orange)
  local O W GR Z wallet pool name slots gpus
  if [ -t 1 ] && [ -z "${NO_COLOR:-}" ]; then O=$'\e[1;38;2;255;77;0m'; W=$'\e[38;2;237;232;219m'; GR=$'\e[38;2;163;157;144m'; Z=$'\e[0m'; else O=""; W=""; GR=""; Z=""; fi
  wallet=$(grep -E "^WALLET_ADDRESS=" .env 2>/dev/null | cut -d= -f2-); pool=$(grep -E "^POOL_URL=" .env 2>/dev/null | cut -d= -f2-)
  name=$(grep -E "^WORKER_NAME=" .env 2>/dev/null | cut -d= -f2-); slots=$(grep -E "^LLAMA_PARALLEL=" .env 2>/dev/null | cut -d= -f2-)
  [ ${#wallet} -gt 20 ] && wallet="${wallet:0:10}…${wallet: -8}"
  gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | sed 's/NVIDIA GeForce //' | sort | uniq -c | awk '{n=$1; $1=""; sub(/^ /,""); printf "%s%d x %s", (c++?" + ":""), n, $0}')
  echo "  ${GR}Pool         ${Z}${W}${pool:-?}${Z}"
  echo "  ${GR}Address      ${Z}${W}${wallet:-?}${Z}"
  echo "  ${GR}Worker       ${Z}${W}${name:-?}  (-g0, -g1 …)${Z}"
  echo "  ${GR}GPUs         ${Z}${W}${gpus:-no GPU seen} · ${slots:-192} windows${Z}"
  echo "  ${GR}Fee          ${Z}${O}1%${Z}"
  echo
}
case "${1:-}" in
  logo)    aria_logo "$ENERGY_VERSION"; aria_fiche ;;
  start)   aria_logo "$ENERGY_VERSION"; aria_fiche; docker compose --env-file .env up -d && echo "started. See: ./view   (raw log: ./tsc logs)" ;;
  stop)
     docker compose --env-file .env down
     # cleanup of orphan GPU processes: if the engine was killed during a
     # CUDA load, a llama-server can outlive the container and hold the
     # VRAM. Only processes owned by the current user are killed.
     if command -v nvidia-smi >/dev/null 2>&1; then
       for pid in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null); do
         [ -z "$pid" ] && continue
         own=$(ps -o user= -p "$pid" 2>/dev/null | tr -d " ")
         cmd=$(ps -o comm= -p "$pid" 2>/dev/null | tr -d " ")
         if [ "$own" = "$(id -un)" ] && [ "$cmd" = "llama-server" ]; then
           echo "  orphan GPU process $pid ($cmd) -> stopping"
           kill "$pid" 2>/dev/null; sleep 2; kill -9 "$pid" 2>/dev/null
         fi
       done
     fi
     echo "stopped."
     ;;
  restart) docker compose --env-file .env down; docker compose --env-file .env up -d ;;
  status)  docker compose --env-file .env ps ;;
  view|vue) shift; exec ./view "$@" ;;                                    # state + compact log (./view check | 10m | fix)
  logs)    docker compose --env-file .env logs -f --tail 100 ${2:-} ;;   # follows continuously (Ctrl-C to quit)
  tail)    docker compose --env-file .env logs --tail ${2:-100} ;;       # prints and returns
  verify)
     IMG=$(grep -E "^BACKEND_IMAGE=" .env | cut -d= -f2-); PRX=$(grep -E "^PROXY_IMAGE=" .env | cut -d= -f2-)
     [ -n "$IMG" ] || { echo "BACKEND_IMAGE missing from .env"; exit 1; }
     echo "Checking $IMG:"
     docker run --rm --entrypoint sh "$IMG" -c 'cat /app/ARIA-ENERGY-1.7 2>/dev/null | sed "s/^/  /"; for f in /usr/local/bin/llama-server /usr/local/lib/libllama.so.0 /usr/local/lib/libggml-cuda.so.0 /usr/local/lib/libcublas.so.12; do [ -e "$f" ] && echo "  ok  $f" || echo "  MISSING $f"; done; /usr/local/bin/llama-server --version 2>&1 | head -1 | sed "s/^/  /"'
     echo "Checking $PRX:"
     docker run --rm --entrypoint sh "$PRX" -c 'for f in /app/src/worker_client.cpython-310-x86_64-linux-gnu.so /app/src/components/vdf_service.cpython-310-x86_64-linux-gnu.so; do [ -e "$f" ] && echo "  ok  $f" || echo "  MISSING $f"; done; [ -e /app/src/worker_client.py ] && echo "  ATTN plain-text worker_client.py present" || echo "  ok  official proxy v1.2.0a + compiled ARIA modules"'
     ;;
  delete)
     read -r -p "Delete everything, including the downloaded model? (type DELETE): " c
     if [ "$c" = "DELETE" ]; then
       docker compose --env-file .env down -v 2>/dev/null
       D="$(pwd)"
       docker run --rm -v "$D:/target" alpine sh -c "rm -rf /target/* /target/.[!.]* 2>/dev/null" 2>/dev/null
       cd .. && rm -rf "$D" 2>/dev/null || sudo rm -rf "$D"
       echo "deleted: $D"
     else echo "cancelled."; fi
     ;;
  *) echo "usage: ./tsc {start|stop|restart|status|view|logs|tail|verify|delete}" ;;
esac
