#!/bin/bash
#
# DGX Spark デュアル構成セットアップスクリプト
# Usage: curl -sL <url> | bash -s -- [command]
#
# Commands:
#   network    - QSFPインターフェースのIP設定
#   ssh        - 対向ノードへのSSH鍵配布
#   docker     - Docker権限設定
#   vllm-pull  - vLLMイメージ取得
#   cluster    - vLLMクラスター起動
#   all        - 全セットアップ実行
#
set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info()  { echo -e "${BLUE}[INFO]${NC} $1"; }
log_ok()    { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# QSFPインターフェース検出
detect_qsfp_interface() {
    local iface
    iface=$(ibdev2netdev 2>/dev/null | grep "(Up)" | awk '{print $5}' | head -1)
    if [[ -z "$iface" ]]; then
        # フォールバック: enp1s0f で始まるインターフェースを探す
        iface=$(ip link show | grep -oP 'enp1s0f[0-9]+np[0-9]+' | head -1)
    fi
    echo "$iface"
}

# 現在のIPアドレス取得
get_current_ip() {
    local iface=$1
    ip -4 addr show "$iface" 2>/dev/null | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | head -1
}

# ネットワーク設定
cmd_network() {
    log_info "QSFPインターフェースを検出中..."
    
    local iface
    iface=$(detect_qsfp_interface)
    
    if [[ -z "$iface" ]]; then
        log_error "QSFPインターフェースが見つかりません"
        log_info "ibdev2netdev の出力を確認してください"
        exit 1
    fi
    
    log_ok "検出: $iface"
    
    # 現在のIP確認
    local current_ip
    current_ip=$(get_current_ip "$iface")
    
    if [[ -n "$current_ip" ]]; then
        log_info "現在のIP: $current_ip"
        read -rp "このIPを使用しますか？ [Y/n]: " use_current
        if [[ "${use_current,,}" != "n" ]]; then
            log_ok "設定完了: $iface = $current_ip"
            return 0
        fi
    fi
    
    # IP入力
    echo ""
    log_info "このノードのIPアドレスを設定します"
    log_info "例: Node 1 = 192.168.100.10, Node 2 = 192.168.100.11"
    read -rp "IPアドレス (例: 192.168.100.10): " new_ip
    read -rp "サブネットマスク [24]: " subnet
    subnet=${subnet:-24}
    
    # 設定適用
    log_info "IPアドレスを設定中..."
    sudo ip addr flush dev "$iface" 2>/dev/null || true
    sudo ip addr add "${new_ip}/${subnet}" dev "$iface"
    sudo ip link set "$iface" up
    
    log_ok "設定完了: $iface = $new_ip/$subnet"
    
    # 永続化確認
    read -rp "netplanで永続化しますか？ [Y/n]: " persist
    if [[ "${persist,,}" != "n" ]]; then
        local netplan_file="/etc/netplan/99-dgx-spark-qsfp.yaml"
        log_info "netplan設定を作成中..."
        
        sudo tee "$netplan_file" > /dev/null << EOF
network:
  version: 2
  ethernets:
    ${iface}:
      addresses:
        - ${new_ip}/${subnet}
EOF
        
        sudo netplan apply
        log_ok "永続化完了: $netplan_file"
    fi
    
    # 環境変数出力
    echo ""
    log_info "以下の環境変数をエクスポートしてください:"
    echo "export MN_IF_NAME=$iface"
    echo "export VLLM_HOST_IP=$new_ip"
}

# SSH鍵配布
cmd_ssh() {
    log_info "SSH鍵の設定を開始します"
    
    # 鍵がなければ生成
    if [[ ! -f ~/.ssh/id_ed25519 ]]; then
        log_info "SSH鍵を生成中..."
        ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519
        log_ok "鍵生成完了"
    else
        log_ok "既存の鍵を使用: ~/.ssh/id_ed25519"
    fi
    
    # 対向ノードのIP入力
    read -rp "対向ノードのIPアドレス: " remote_ip
    read -rp "対向ノードのユーザー名 [$USER]: " remote_user
    remote_user=${remote_user:-$USER}
    
    log_info "対向ノードに公開鍵を配布中..."
    log_warn "パスワードを求められます"
    
    ssh-copy-id "${remote_user}@${remote_ip}"
    
    log_ok "SSH鍵配布完了"
    
    # 疎通確認
    log_info "接続テスト中..."
    if ssh -o BatchMode=yes -o ConnectTimeout=5 "${remote_user}@${remote_ip}" "echo OK" &>/dev/null; then
        log_ok "パスワードなしSSH接続成功"
    else
        log_error "接続に失敗しました"
        exit 1
    fi
}

# Docker権限設定
cmd_docker() {
    log_info "Docker権限を設定中..."
    
    if groups | grep -q docker; then
        log_ok "既にdockerグループに所属しています"
    else
        sudo groupadd docker 2>/dev/null || true
        sudo usermod -aG docker "$USER"
        log_ok "dockerグループに追加しました"
        log_warn "変更を反映するには再ログインするか 'newgrp docker' を実行してください"
    fi
    
    # NVIDIA Container Toolkit確認
    if docker run --rm --gpus all nvidia/cuda:13.0.1-devel-ubuntu24.04 nvidia-smi &>/dev/null; then
        log_ok "NVIDIA Container Toolkit正常動作"
    else
        log_error "GPU付きコンテナが起動できません"
        log_info "NVIDIA Container Toolkitをインストールしてください"
        exit 1
    fi
}

# vLLMイメージ取得
cmd_vllm_pull() {
    local image="nvcr.io/nvidia/vllm:25.11-py3"
    
    log_info "vLLMイメージを取得中..."
    log_info "Image: $image"
    
    docker pull "$image"
    
    log_ok "取得完了"
    echo ""
    echo "export VLLM_IMAGE=$image"
}

# vLLMクラスター起動
cmd_cluster() {
    log_info "vLLMクラスターを起動します"
    
    # 環境変数確認
    if [[ -z "${VLLM_IMAGE:-}" ]]; then
        export VLLM_IMAGE="nvcr.io/nvidia/vllm:25.11-py3"
        log_warn "VLLM_IMAGE未設定、デフォルト使用: $VLLM_IMAGE"
    fi
    
    if [[ -z "${VLLM_HOST_IP:-}" ]]; then
        local iface
        iface=$(detect_qsfp_interface)
        if [[ -n "$iface" ]]; then
            VLLM_HOST_IP=$(get_current_ip "$iface")
        fi
        
        if [[ -z "${VLLM_HOST_IP:-}" ]]; then
            read -rp "このノードのクラスター通信用IP: " VLLM_HOST_IP
        fi
    fi
    
    if [[ -z "${MN_IF_NAME:-}" ]]; then
        MN_IF_NAME=$(detect_qsfp_interface)
    fi
    
    log_info "設定:"
    echo "  VLLM_IMAGE:   $VLLM_IMAGE"
    echo "  VLLM_HOST_IP: $VLLM_HOST_IP"
    echo "  MN_IF_NAME:   $MN_IF_NAME"
    echo ""
    
    # ノードタイプ選択
    echo "このノードの役割を選択してください:"
    echo "  1) ヘッドノード (Node 1)"
    echo "  2) ワーカーノード (Node 2)"
    read -rp "選択 [1/2]: " node_type
    
    # run_cluster.sh取得
    if [[ ! -f ./run_cluster.sh ]]; then
        log_info "run_cluster.sh をダウンロード中..."
        wget -q https://raw.githubusercontent.com/vllm-project/vllm/refs/heads/main/examples/online_serving/run_cluster.sh
        chmod +x run_cluster.sh
    fi
    
    local head_ip
    if [[ "$node_type" == "1" ]]; then
        head_ip="$VLLM_HOST_IP"
        log_info "ヘッドノードとして起動中..."
        
        bash run_cluster.sh "$VLLM_IMAGE" "$head_ip" --head ~/.cache/huggingface \
            -e VLLM_HOST_IP="$VLLM_HOST_IP" \
            -e UCX_NET_DEVICES="$MN_IF_NAME" \
            -e NCCL_SOCKET_IFNAME="$MN_IF_NAME" \
            -e OMPI_MCA_btl_tcp_if_include="$MN_IF_NAME" \
            -e GLOO_SOCKET_IFNAME="$MN_IF_NAME" \
            -e TP_SOCKET_IFNAME="$MN_IF_NAME" \
            -e RAY_memory_monitor_refresh_ms=0 \
            -e MASTER_ADDR="$head_ip"
    else
        read -rp "ヘッドノード(Node 1)のIP: " head_ip
        log_info "ワーカーノードとして起動中..."
        
        bash run_cluster.sh "$VLLM_IMAGE" "$head_ip" --worker ~/.cache/huggingface \
            -e VLLM_HOST_IP="$VLLM_HOST_IP" \
            -e UCX_NET_DEVICES="$MN_IF_NAME" \
            -e NCCL_SOCKET_IFNAME="$MN_IF_NAME" \
            -e OMPI_MCA_btl_tcp_if_include="$MN_IF_NAME" \
            -e GLOO_SOCKET_IFNAME="$MN_IF_NAME" \
            -e TP_SOCKET_IFNAME="$MN_IF_NAME" \
            -e RAY_memory_monitor_refresh_ms=0 \
            -e MASTER_ADDR="$head_ip"
    fi
}

# 全セットアップ
cmd_all() {
    log_info "=== DGX Spark デュアル構成 フルセットアップ ==="
    echo ""
    
    cmd_docker
    echo ""
    
    cmd_network
    echo ""
    
    read -rp "対向ノードへのSSH鍵配布を行いますか？ [Y/n]: " do_ssh
    if [[ "${do_ssh,,}" != "n" ]]; then
        cmd_ssh
        echo ""
    fi
    
    cmd_vllm_pull
    echo ""
    
    log_ok "セットアップ完了！"
    echo ""
    log_info "次のステップ:"
    echo "  1. 対向ノードでも同じスクリプトを実行"
    echo "  2. 両ノードで 'dgx-spark-setup.sh cluster' を実行"
    echo "  3. ヘッドノードでモデルを起動"
}

# ヘルプ
cmd_help() {
    cat << 'EOF'
DGX Spark デュアル構成セットアップスクリプト

Usage:
  dgx-spark-setup.sh <command>

Commands:
  network    QSFPインターフェースのIP設定
  ssh        対向ノードへのSSH鍵配布
  docker     Docker権限設定
  vllm-pull  vLLMイメージ取得
  cluster    vLLMクラスター起動
  all        全セットアップ実行（推奨）

Examples:
  # フルセットアップ
  ./dgx-spark-setup.sh all

  # ネットワークのみ設定
  ./dgx-spark-setup.sh network

  # クラスター起動
  export VLLM_HOST_IP=192.168.100.10
  export MN_IF_NAME=enp1s0f1np1
  ./dgx-spark-setup.sh cluster

ワンライナー実行:
  curl -sL https://example.com/dgx-spark-setup.sh | bash -s -- all
EOF
}

# メイン
main() {
    local cmd="${1:-help}"
    
    case "$cmd" in
        network)   cmd_network ;;
        ssh)       cmd_ssh ;;
        docker)    cmd_docker ;;
        vllm-pull) cmd_vllm_pull ;;
        cluster)   cmd_cluster ;;
        all)       cmd_all ;;
        help|--help|-h) cmd_help ;;
        *)
            log_error "Unknown command: $cmd"
            cmd_help
            exit 1
            ;;
    esac
}

main "$@"