#!/usr/bin/env bash
# RigHost — Vast.ai Compatibility Checker
# Usage: curl -sSL https://www.righost.com/check.sh | bash
# Source: https://www.righost.com/check.sh
set -euo pipefail

API="https://www.righost.com/api/check"
BOLD="\033[1m"
RESET="\033[0m"
GREEN="\033[32m"
YELLOW="\033[33m"
RED="\033[31m"
BLUE="\033[34m"
GRAY="\033[90m"

echo ""
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo -e "${BOLD}  RigHost — Vast.ai Compatibility Checker${RESET}"
echo -e "${GRAY}  https://www.righost.com${RESET}"
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""
echo -e "${GRAY}Detecting your hardware — nothing is stored or transmitted until the check runs.${RESET}"
echo ""

# ── GPU detection ──────────────────────────────────
GPU_NAME=""
GPU_VRAM=0
if command -v nvidia-smi &>/dev/null; then
    GPU_RAW=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || true)
    if [[ -n "$GPU_RAW" ]]; then
        GPU_NAME=$(echo "$GPU_RAW" | cut -d',' -f1 | xargs)
        GPU_VRAM_MB=$(echo "$GPU_RAW" | cut -d',' -f2 | xargs)
        GPU_VRAM=$(( (GPU_VRAM_MB + 512) / 1024 ))  # MB → GB, rounded
    fi
fi
if [[ -z "$GPU_NAME" ]]; then
    GPU_NAME="unknown"
    echo -e "  ${YELLOW}⚠  GPU:${RESET}     nvidia-smi not found or no GPU detected"
    read -rp "     Enter GPU model manually (e.g. RTX 3090): " GPU_NAME_INPUT
    GPU_NAME="${GPU_NAME_INPUT:-unknown}"
    read -rp "     Enter GPU VRAM in GB (e.g. 24): " GPU_VRAM_INPUT
    GPU_VRAM="${GPU_VRAM_INPUT:-0}"
else
    echo -e "  ${GREEN}✓  GPU:${RESET}     ${GPU_NAME} (${GPU_VRAM} GB VRAM)"
fi

# ── CPU ────────────────────────────────────────────
CORES=$(nproc 2>/dev/null || grep -c ^processor /proc/cpuinfo 2>/dev/null || echo 0)
echo -e "  ${GREEN}✓  CPU:${RESET}     ${CORES} logical cores"

# ── RAM ────────────────────────────────────────────
RAM_MB=$(free -m 2>/dev/null | awk '/^Mem:/{print $2}' || echo 0)
RAM_GB=$(awk "BEGIN {printf \"%.0f\", ${RAM_MB}/1024}")
echo -e "  ${GREEN}✓  RAM:${RESET}     ~${RAM_GB} GB"

# ── Storage (available on /) ───────────────────────
STORAGE_GB=$(df -BG / 2>/dev/null | awk 'NR==2{gsub(/G/,"",$4); print $4}' || echo 0)
echo -e "  ${GREEN}✓  Storage:${RESET} ${STORAGE_GB} GB available on /"

# ── OS ─────────────────────────────────────────────
OS_VALUE="linux-other"
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    OS_ID="${ID:-linux}"
    OS_VER="${VERSION_ID:-}"
    if [[ "$OS_ID" == "ubuntu" && "$OS_VER" == "22.04" ]]; then OS_VALUE="ubuntu22"
    elif [[ "$OS_ID" == "ubuntu" && "$OS_VER" == "20.04" ]]; then OS_VALUE="ubuntu20"
    elif [[ "$OS_ID" == "ubuntu" ]]; then OS_VALUE="linux-other"
    fi
    echo -e "  ${GREEN}✓  OS:${RESET}      ${PRETTY_NAME:-$OS_ID $OS_VER}"
else
    echo -e "  ${YELLOW}⚠  OS:${RESET}      Could not detect — assuming linux-other"
fi

# ── Docker ────────────────────────────────────────
DOCKER_VALUE="no"
if command -v docker &>/dev/null && docker --version &>/dev/null 2>&1; then
    DOCKER_VERSION=$(docker --version 2>/dev/null | awk '{print $3}' | tr -d ',')
    DOCKER_VALUE="yes"
    echo -e "  ${GREEN}✓  Docker:${RESET}  Installed (${DOCKER_VERSION})"
else
    echo -e "  ${YELLOW}⚠  Docker:${RESET}  Not detected"
fi

# ── Network speeds ─────────────────────────────────
echo ""
echo -e "${BOLD}Network speeds (run a speed test if you don't know these):${RESET}"
echo -e "${GRAY}  Tip: visit https://www.speedtest.net or run: curl -s https://raw.githubusercontent.com/sivel/speedtest-cli/master/speedtest.py | python3${RESET}"
echo ""
read -rp "  Upload speed in Mbps (e.g. 500): " UPLOAD_INPUT
UPLOAD="${UPLOAD_INPUT:-0}"
read -rp "  Download speed in Mbps (e.g. 500): " DOWNLOAD_INPUT
DOWNLOAD="${DOWNLOAD_INPUT:-0}"

# ── Build API URL & call ──────────────────────────
GPU_ENC=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${GPU_NAME}'))" 2>/dev/null \
          || echo "${GPU_NAME// /+}")

echo ""
echo -e "${GRAY}Running compatibility check…${RESET}"

RESPONSE=$(curl -sf \
    "${API}?gpu=${GPU_ENC}&vram=${GPU_VRAM}&upload=${UPLOAD}&download=${DOWNLOAD}&storage=${STORAGE_GB}&os=${OS_VALUE}&docker=${DOCKER_VALUE}&cores=${CORES}&ram=${RAM_GB}" \
    2>/dev/null || true)

if [[ -z "$RESPONSE" ]]; then
    echo -e "${RED}Error: Could not reach the RigHost API. Check your internet connection.${RESET}"
    exit 1
fi

# ── Parse & display results (pure bash+awk, no jq) ──
SCORE=$(echo "$RESPONSE" | grep -o '"score": *[0-9]*' | grep -o '[0-9]*' | head -1)
VERDICT=$(echo "$RESPONSE" | grep -o '"verdict": *"[^"]*"' | grep -o '"[^"]*"$' | tr -d '"' | head -1)

echo ""
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo -e "${BOLD}  YOUR RESULTS${RESET}"
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""

# Score badge
case "$VERDICT" in
    ready)      echo -e "  ${BOLD}${GREEN}✓ READY TO HOST${RESET}  —  ${BOLD}${SCORE}/100 points${RESET}" ;;
    needs-work) echo -e "  ${BOLD}${YELLOW}⚠ NEEDS SOME WORK${RESET}  —  ${BOLD}${SCORE}/100 points${RESET}" ;;
    *)          echo -e "  ${BOLD}${RED}✗ NOT READY YET${RESET}  —  ${BOLD}${SCORE}/100 points${RESET}" ;;
esac
echo ""

# Category breakdown
echo -e "  ${BOLD}Category Breakdown:${RESET}"
echo ""

# Parse category lines from JSON
while IFS= read -r line; do
    NAME=$(echo "$line" | grep -o '"name": *"[^"]*"' | grep -o '"[^"]*"$' | tr -d '"')
    PTS=$(echo "$line" | grep -o '"pts": *[0-9]*' | grep -o '[0-9]*')
    MAX=$(echo "$line" | grep -o '"max": *[0-9]*' | grep -o '[0-9]*')
    STATUS=$(echo "$line" | grep -o '"status": *"[^"]*"' | grep -o '"[^"]*"$' | tr -d '"')
    DETAIL=$(echo "$line" | grep -o '"detail": *"[^"]*"' | grep -o '"[^"]*"$' | tr -d '"')

    [[ -z "$NAME" ]] && continue

    case "$STATUS" in
        pass) ICON="${GREEN}✓${RESET}" ;;
        warn) ICON="${YELLOW}⚠${RESET}" ;;
        *)    ICON="${RED}✗${RESET}" ;;
    esac

    printf "  %b  %-20s %s/%s pts  %s\n" "$ICON" "$NAME" "$PTS" "$MAX" "${GRAY}${DETAIL}${RESET}"
done < <(echo "$RESPONSE" | grep -o '{"name":[^}]*}')

# Recommendations
RECS_RAW=$(echo "$RESPONSE" | grep -o '"recommendations": *\[[^]]*\]' | grep -o '\[[^]]*\]' | tr -d '[]"' | tr ',' '\n')
if [[ -n "$RECS_RAW" ]]; then
    echo ""
    echo -e "  ${BOLD}Recommendations:${RESET}"
    while IFS= read -r rec; do
        rec=$(echo "$rec" | xargs)
        [[ -z "$rec" ]] && continue
        echo -e "  ${YELLOW}→${RESET}  ${rec}"
    done <<< "$RECS_RAW"
fi

echo ""
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo -e "  Start hosting at: ${BLUE}https://vast.ai/host${RESET}"
echo -e "  More info:        ${BLUE}https://www.righost.com${RESET}"
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""