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

API="https://www.righost.com/api/runpod-check.php"
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 — RunPod Provider Compatibility Check${RESET}"
echo -e "${GRAY}  https://www.righost.com/runpod.html${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
GPU_COUNT=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 || true)
    if [[ -n "$GPU_RAW" ]]; then
        GPU_COUNT=$(echo "$GPU_RAW" | wc -l | xargs)
        FIRST=$(echo "$GPU_RAW" | head -1)
        GPU_NAME=$(echo "$FIRST" | cut -d',' -f1 | xargs)
        GPU_VRAM_MB=$(echo "$FIRST" | cut -d',' -f2 | xargs)
        GPU_VRAM=$(( (GPU_VRAM_MB + 512) / 1024 ))
    fi
fi
if [[ -z "$GPU_NAME" ]]; then
    echo -e "  ${YELLOW}⚠  GPU:${RESET}        nvidia-smi not found or no GPU detected"
    read -rp "     Enter GPU model manually (e.g. RTX 4090): " 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}"
    read -rp "     Number of GPUs (default 1): " GPU_COUNT_INPUT
    GPU_COUNT="${GPU_COUNT_INPUT:-1}"
else
    echo -e "  ${GREEN}✓  GPU:${RESET}        ${GPU_NAME} x${GPU_COUNT} (${GPU_VRAM} GB VRAM each)"
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 type + free space (NVMe detection via lsblk) ─
STORAGE_TYPE="nvme"
STORAGE_GB=0
ROOT_DEV=""

# Find device backing / (strip partition suffix to get parent)
if command -v findmnt &>/dev/null; then
    ROOT_SRC=$(findmnt -n -o SOURCE / 2>/dev/null || echo "")
    if [[ -n "$ROOT_SRC" && "$ROOT_SRC" == /dev/* ]]; then
        ROOT_DEV=$(basename "$ROOT_SRC" | sed -E 's/p?[0-9]+$//')
    fi
fi

if command -v lsblk &>/dev/null; then
    if [[ -n "$ROOT_DEV" ]]; then
        ROTA=$(lsblk -dn -o ROTA "/dev/${ROOT_DEV}" 2>/dev/null || echo "1")
        if [[ "$ROOT_DEV" == nvme* ]]; then
            STORAGE_TYPE="nvme"
        elif [[ "$ROTA" == "0" ]]; then
            STORAGE_TYPE="ssd"
        else
            STORAGE_TYPE="hdd"
        fi
    fi
fi

STORAGE_GB=$(df -BG / 2>/dev/null | awk 'NR==2{gsub(/G/,"",$4); print $4}' || echo 0)
case "$STORAGE_TYPE" in
    nvme) echo -e "  ${GREEN}✓  Storage:${RESET}    NVMe, ${STORAGE_GB} GB available on /" ;;
    ssd)  echo -e "  ${YELLOW}⚠  Storage:${RESET}    SATA SSD, ${STORAGE_GB} GB available (NVMe preferred)" ;;
    hdd)  echo -e "  ${RED}✗  Storage:${RESET}    HDD, ${STORAGE_GB} GB available (RunPod requires NVMe)" ;;
    *)    echo -e "  ${YELLOW}⚠  Storage:${RESET}    Unknown type, ${STORAGE_GB} GB available" ;;
esac

# ── OS ─────────────────────────────────────────────
OS_PRETTY="Unknown"
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    OS_PRETTY="${PRETTY_NAME:-$ID $VERSION_ID}"
    echo -e "  ${GREEN}✓  OS:${RESET}         ${OS_PRETTY}"
else
    echo -e "  ${YELLOW}⚠  OS:${RESET}         Could not detect"
fi

# ── Docker + NVIDIA Container Toolkit ──────────────
DOCKER_NVIDIA_VALUE="no"
DOCKER_INSTALLED="no"
NVIDIA_TOOLKIT="no"
if command -v docker &>/dev/null && docker --version &>/dev/null 2>&1; then
    DOCKER_INSTALLED="yes"
    DOCKER_VERSION=$(docker --version 2>/dev/null | awk '{print $3}' | tr -d ',')
fi

if command -v nvidia-container-cli &>/dev/null || command -v nvidia-container-runtime &>/dev/null; then
    NVIDIA_TOOLKIT="yes"
fi

# Alternate detection: package presence
if [[ "$NVIDIA_TOOLKIT" == "no" ]]; then
    if dpkg -l 2>/dev/null | grep -qE 'nvidia-container-toolkit|nvidia-docker2'; then
        NVIDIA_TOOLKIT="yes"
    elif rpm -qa 2>/dev/null | grep -qE 'nvidia-container-toolkit|nvidia-docker2'; then
        NVIDIA_TOOLKIT="yes"
    fi
fi

if [[ "$DOCKER_INSTALLED" == "yes" && "$NVIDIA_TOOLKIT" == "yes" ]]; then
    DOCKER_NVIDIA_VALUE="yes"
    echo -e "  ${GREEN}✓  Docker:${RESET}     ${DOCKER_VERSION} + nvidia-container-toolkit"
elif [[ "$DOCKER_INSTALLED" == "yes" ]]; then
    echo -e "  ${YELLOW}⚠  Docker:${RESET}     ${DOCKER_VERSION} installed, but nvidia-container-toolkit NOT detected"
else
    echo -e "  ${RED}✗  Docker:${RESET}     Not installed"
fi

# ── Static IP / reachability ───────────────────────
PUBLIC_IP=$(curl -s --max-time 4 https://api.ipify.org 2>/dev/null || echo "")
LOCAL_IPS=$(ip -4 addr show scope global 2>/dev/null | awk '/inet /{print $2}' | cut -d/ -f1 | tr '\n' ' ' || echo "")

STATIC_IP_VALUE="unknown"
if [[ -n "$PUBLIC_IP" && -n "$LOCAL_IPS" ]]; then
    if echo "$LOCAL_IPS" | grep -qw "$PUBLIC_IP"; then
        STATIC_IP_VALUE="yes"
        echo -e "  ${GREEN}✓  IPv4:${RESET}       ${PUBLIC_IP} (appears directly routable)"
    else
        STATIC_IP_VALUE="no"
        echo -e "  ${YELLOW}⚠  IPv4:${RESET}       ${PUBLIC_IP} public, local=${LOCAL_IPS% } (likely NAT)"
    fi
else
    echo -e "  ${YELLOW}⚠  IPv4:${RESET}       Could not determine"
fi

# ── Prompts: bandwidth + uptime ────────────────────
echo ""
echo -e "${BOLD}A few things we can't auto-detect:${RESET}"
echo ""

echo -e "${GRAY}  Tip for bandwidth: visit https://www.speedtest.net or run:${RESET}"
echo -e "${GRAY}    pip install speedtest-cli && speedtest-cli${RESET}"
read -rp "  Symmetric bandwidth in Mbps (e.g. 1000): " BW_INPUT
BANDWIDTH="${BW_INPUT:-0}"

read -rp "  Do you commit to running 24/7? [y/N]: " UP_INPUT
case "${UP_INPUT:-n}" in
    y|Y|yes|YES) UPTIME_TIER="24_7" ;;
    *)           UPTIME_TIER="occasional" ;;
esac

if [[ "$STATIC_IP_VALUE" == "unknown" ]]; then
    read -rp "  Does your IP appear static with open ports? [y/N/unknown]: " IP_INPUT
    case "${IP_INPUT:-unknown}" in
        y|Y|yes|YES) STATIC_IP_VALUE="yes" ;;
        n|N|no|NO)   STATIC_IP_VALUE="no"  ;;
        *)           STATIC_IP_VALUE="unknown" ;;
    esac
fi

# ── Build JSON payload ─────────────────────────────
echo ""
echo -e "${GRAY}Running RunPod compatibility check…${RESET}"

JSON_PAYLOAD=$(cat <<JSONEOF
{
  "gpu_model": "${GPU_NAME}",
  "gpu_count": ${GPU_COUNT},
  "vram_gb": ${GPU_VRAM},
  "bandwidth_mbps": ${BANDWIDTH},
  "static_ip": "${STATIC_IP_VALUE}",
  "nvme_gb": ${STORAGE_GB},
  "storage_type": "${STORAGE_TYPE}",
  "uptime_tier": "${UPTIME_TIER}",
  "docker_nvidia": "${DOCKER_NVIDIA_VALUE}",
  "cpu_cores": ${CORES},
  "ram_gb": ${RAM_GB}
}
JSONEOF
)

RESPONSE=$(curl -sf -X POST "${API}" \
    -H "Content-Type: application/json" \
    -d "$JSON_PAYLOAD" 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 (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 RUNPOD RESULTS${RESET}"
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""

case "$VERDICT" in
    ready)      echo -e "  ${BOLD}${GREEN}✓ READY TO APPLY${RESET}    —  ${BOLD}${SCORE}/100 points${RESET}" ;;
    near-ready) 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 ""

echo -e "  ${BOLD}Category Breakdown:${RESET}"
echo ""
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  %-34s %s/%s pts  %s\n" "$ICON" "$NAME" "$PTS" "$MAX" "${GRAY}${DETAIL}${RESET}"
done < <(echo "$RESPONSE" | grep -o '{"name":[^}]*}')

RECS_RAW=$(echo "$RESPONSE" | grep -o '"recommendations": *\[[^]]*\]' | grep -o '\[[^]]*\]' | sed 's/^\[//;s/\]$//' | sed 's/","/\n/g' | tr -d '"')
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 "  Apply as a provider: ${BLUE}https://www.runpod.io/console/user/settings${RESET}  ${GRAY}(verify official URL)${RESET}"
echo -e "  More info:           ${BLUE}https://www.righost.com/runpod.html${RESET}"
echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""

# Exit 0 if ready, else 1
[[ "$VERDICT" == "ready" ]] && exit 0 || exit 1
