#!/usr/bin/env bash
# Heat GPU Cloud — Provider Client インストーラ
# Usage:  curl -sSL https://heatgpu.com/provider/install.sh | bash
#         または トークンを直接渡す:
#         curl -sSL https://heatgpu.com/provider/install.sh | HEAT_PROVIDER_TOKEN=xxx bash
set -euo pipefail

VERSION="0.1.0"
INSTALL_DIR="${HEAT_PROVIDER_HOME:-$HOME/.heat-provider}"
SCRIPT_URL="https://heatgpu.com/provider/heat_provider.py"

c_neon='\033[38;2;0;255;163m'
c_gold='\033[38;2;255;213;107m'
c_mute='\033[38;5;245m'
c_red='\033[38;5;203m'
c_off='\033[0m'
b='\033[1m'

echo
echo -e "${c_neon}${b}─────────────────────────────────────────────${c_off}"
echo -e "${c_neon}${b}  Heat GPU Cloud — Provider Client v${VERSION}${c_off}"
echo -e "${c_neon}${b}─────────────────────────────────────────────${c_off}"
echo

# Python 3.8+ チェック
if ! command -v python3 >/dev/null 2>&1; then
  echo -e "${c_red}✕ Python 3 が必要です。https://www.python.org/ からインストールしてください。${c_off}"
  exit 1
fi
PY_VER=$(python3 -c 'import sys; print("%d.%d" % sys.version_info[:2])')
echo -e "  ${c_mute}python:${c_off}  $(which python3)  ${c_mute}(${PY_VER})${c_off}"

# nvidia-smi 警告 (Mac でも 警告のみで継続)
if ! command -v nvidia-smi >/dev/null 2>&1; then
  echo -e "  ${c_gold}!  nvidia-smi が見つかりません。NVIDIA GPU + ドライバが必要です。${c_off}"
  echo -e "     ${c_mute}インストールは続行しますが、起動時にエラーになります。${c_off}"
else
  GPU=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
  echo -e "  ${c_mute}gpu:${c_off}     ${GPU}"
fi

mkdir -p "$INSTALL_DIR"
echo -e "  ${c_mute}install:${c_off} ${INSTALL_DIR}"

# スクリプト取得
echo -e "  ${c_mute}fetch:${c_off}   ${SCRIPT_URL}"
if command -v curl >/dev/null 2>&1; then
  curl -fsSL "$SCRIPT_URL" -o "$INSTALL_DIR/heat_provider.py"
elif command -v wget >/dev/null 2>&1; then
  wget -qO "$INSTALL_DIR/heat_provider.py" "$SCRIPT_URL"
else
  echo -e "${c_red}✕ curl or wget が必要です${c_off}"
  exit 1
fi
chmod +x "$INSTALL_DIR/heat_provider.py"

# 起動スクリプト
cat > "$INSTALL_DIR/start.sh" <<EOF
#!/usr/bin/env bash
# Heat Provider Client 起動スクリプト
# 環境変数 HEAT_PROVIDER_TOKEN を設定してください
cd "${INSTALL_DIR}"
exec python3 heat_provider.py "\$@"
EOF
chmod +x "$INSTALL_DIR/start.sh"

# PATH 追加 (~/.zshrc or ~/.bashrc)
SHELL_RC=""
if [[ -n "${ZSH_VERSION:-}" ]] || [[ "${SHELL:-}" == */zsh ]]; then
  SHELL_RC="$HOME/.zshrc"
elif [[ -n "${BASH_VERSION:-}" ]] || [[ "${SHELL:-}" == */bash ]]; then
  SHELL_RC="$HOME/.bashrc"
fi
if [[ -n "$SHELL_RC" ]] && ! grep -q "heat-provider" "$SHELL_RC" 2>/dev/null; then
  cat >> "$SHELL_RC" <<EOF

# Heat GPU Cloud Provider Client
alias heat-provider="${INSTALL_DIR}/start.sh"
EOF
  echo -e "  ${c_mute}alias:${c_off}   heat-provider → ${INSTALL_DIR}/start.sh"
fi

echo
echo -e "${c_neon}${b}  ✓ インストール完了${c_off}"
echo
echo -e "  ${b}次のステップ${c_off}"
echo -e "  ${c_mute}─────────────${c_off}"
echo
echo -e "  1. Client Token を取得 (Provider Portal):"
echo -e "     ${c_neon}https://heatgpu.com/provider/gpus.html${c_off}"
echo
echo -e "  2. 接続開始:"
echo -e "     ${c_gold}heat-provider --token YOUR_CLIENT_TOKEN${c_off}"
echo
echo -e "     または環境変数:"
echo -e "     ${c_gold}export HEAT_PROVIDER_TOKEN=YOUR_TOKEN${c_off}"
echo -e "     ${c_gold}heat-provider${c_off}"
echo
echo -e "  3. systemd サービス化 (Linux・常時稼働):"
echo -e "     ${c_neon}sudo curl -fsSL https://heatgpu.com/provider/heat-provider.service \\${c_off}"
echo -e "     ${c_neon}     -o /etc/systemd/system/heat-provider.service${c_off}"
echo -e "     ${c_neon}sudo systemctl enable --now heat-provider${c_off}"
echo

# トークンが環境変数で渡されていれば即起動
if [[ -n "${HEAT_PROVIDER_TOKEN:-}" ]]; then
  echo -e "${c_gold}  HEAT_PROVIDER_TOKEN が設定されています。今すぐ接続を開始します...${c_off}"
  echo
  exec python3 "$INSTALL_DIR/heat_provider.py" --token "$HEAT_PROVIDER_TOKEN"
fi