#!/bin/bash
# ==================== KOGUN 贡献者节点 — 跨平台安装器 v3.1 ====================
# 一条命令搞定所有系统的一键部署
# 
# 用法:
#   bash <(curl -sL https://kogun.ai/install.sh)
#   bash install.sh --auto           # 全自动安装（不提问）
#   bash install.sh --wsl2           # Windows 下强制 WSL2 安装
#   bash install.sh --vllm           # 只部署 vLLM（已有 CLI）
# ==========================================================================

set -euo pipefail

# ---- 颜色 ----
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; NC='\033[0m'

log()     { echo -e "${GREEN}✅${NC} $1"; }
log_info(){ echo -e "${BLUE}ℹ️${NC} $1"; }
log_warn(){ echo -e "${YELLOW}⚠️${NC} $1"; }
log_err() { echo -e "${RED}❌${NC} $1"; }
step()    { echo -e "\n${CYAN}━━━ $1 ━━━${NC}"; }

# ---- 全局配置 ----
KOGUN_API_URL="${KOGUN_API_URL:-https://api.kogun.cn}"
AUTO_MODE=false
FORCE_WSL2=false
ONLY_VLLM=false
SKIP_WSL2=false

# ---- 解析参数 ----
for arg in "$@"; do
  case "$arg" in
    --auto)   AUTO_MODE=true ;;
    --wsl2)   FORCE_WSL2=true ;;
    --vllm)   ONLY_VLLM=true ;;
    --no-wsl) SKIP_WSL2=true ;;
    --help|-h)
      echo "KOGUN 安装器 v3.1"
      echo ""
      echo "用法: bash install.sh [选项]"
      echo ""
      echo "选项:"
      echo "  --auto         全自动（不提问，默认选择）"
      echo "  --wsl2         Windows 下强制安装 WSL2"
      echo "  --vllm         只部署 vLLM（已有 CLI）"
      echo "  --no-wsl       Windows 下跳过 WSL2"
      echo "  --help         显示本帮助"
      exit 0
      ;;
  esac
done

# ---- Banner ----
echo ""
echo -e "${BLUE}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║${NC}  ${CYAN}KOGUN 贡献者节点 — 一键安装${NC}                         ${BLUE}║${NC}"
echo -e "${BLUE}║${NC}  支持 Linux / macOS / Windows (WSL2)                 ${BLUE}║${NC}"
echo -e "${BLUE}║${NC}  有任务才跑 | 按 token 计费 | 70% 分润                ${BLUE}║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════════════════╝${NC}"
echo ""

KOGUN_HOME="$HOME/.kogun"
mkdir -p "$KOGUN_HOME"/{bin,logs,config,data}

# ============================================================
# PHASE 1: 系统检测
# ============================================================
step "1/5 系统环境检测"

ARCH=$(uname -m)
RAW_OS=$(uname -s)

case "$RAW_OS" in
  Linux)                        OS="linux"   ;;
  Darwin)                       OS="macos"   ;;
  MINGW*|MSYS*|CYGWIN*)         OS="windows" ;;
  *)                            OS="$RAW_OS" ;;
esac

echo "  🖥️  系统: $OS ($RAW_OS) | 架构: $ARCH"

# 如果是 WSL2 内部
IS_WSL=false
if [[ -f /proc/sys/fs/binfmt_misc/WSLInterop ]] || [[ -n "${WSL_DISTRO_NAME:-}" ]]; then
  IS_WSL=true
  echo -e "  ${GREEN}检测到 WSL2 环境${NC}"
fi

# 检测 Python
PYTHON=""
if command -v python3 &>/dev/null; then
  PYTHON="python3"
elif command -v python &>/dev/null && python --version 2>&1 | grep -q "^Python 3"; then
  PYTHON="python"
fi

if [[ -z "$PYTHON" ]]; then
  if [[ "$OS" == "linux" ]] && command -v apt &>/dev/null; then
    log_info "安装 Python 3..."
    apt-get install -y -qq python3 python3-pip python3-venv 2>/dev/null || true
    PYTHON="python3"
  else
    log_err "未检测到 Python 3，请手动安装:"
    echo "  Linux:  sudo apt install python3 python3-pip"
    echo "  macOS:  brew install python"
    echo "  Windows: https://python.org 下载安装"
    exit 1
  fi
fi

echo "  🐍 Python: $($PYTHON --version 2>&1)"

# 检测 NVIDIA GPU
HAS_NVIDIA=false
GPU_NAME=""
GPU_VRAM_MB=0

if command -v nvidia-smi &>/dev/null; then
  GPU_LINE=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null | head -1)
  if [[ -n "$GPU_LINE" ]]; then
    GPU_NAME=$(echo "$GPU_LINE" | cut -d, -f1 | xargs)
    VRAM_STR=$(echo "$GPU_LINE" | cut -d, -f2 | xargs)
    GPU_VRAM_MB=${VRAM_STR% MiB}
    HAS_NVIDIA=true
    echo -e "  🎮 GPU: ${CYAN}$GPU_NAME${NC} ($GPU_VRAM_MB MB)"
  fi
fi

echo "  📂 安装目录: $KOGUN_HOME"

# ============================================================
# PHASE 2: Windows → WSL2 自动安装（一键搞定）
# ============================================================
if [[ "$OS" == "windows" ]] && ! $IS_WSL; then
  step "2/5 Windows WSL2 集成"

  # 如果是 --auto 或 --wsl2 或用户同意，走 WSL2 全自动安装
  DO_WSL=false
  if $FORCE_WSL2 || $AUTO_MODE; then
    DO_WSL=true
  elif ! $SKIP_WSL2; then
    echo -e "${CYAN}Windows 原生不支持 vLLM 推理${NC}"
    echo "  KOGUN 可以自动安装 WSL2 + Ubuntu + CUDA + vLLM"
    echo "  装完后 Windows 也能跑 GPU 推理！"
    echo ""
    read -rp "一键安装 WSL2 环境? [Y/n]: " ans
    if [[ ! "$ans" =~ ^[Nn] ]]; then
      DO_WSL=true
    fi
  fi

  if $DO_WSL; then
    log_info "正在安装 WSL2 + Ubuntu..."

    # ─── 3a. 启用 WSL 功能 ───
    if ! wsl.exe --status &>/dev/null; then
      log_info "启用 Windows 功能: WSL..."
      powershell.exe -Command "Enable-WindowsOptionalFeature -Online -FeatureName Microsoft-Windows-Subsystem-Linux -All -NoRestart" 2>/dev/null || true
    fi

    # ─── 3b. 设置 WSL2 默认版本 ───
    WSL_DEF=$(wsl.exe --status 2>/dev/null | grep "Default Version" | grep -oP '\d+' || echo "1")
    if [[ "$WSL_DEF" == "1" ]]; then
      log_info "设置 WSL2 为默认..."
      wsl.exe --set-default-version 2 2>/dev/null || log_warn "需手动启用 WSL2"
    fi

    # ─── 3c. 安装/复用 Ubuntu 发行版 ───
    UBUNTU_NAME=$(wsl.exe -l -q 2>/dev/null | grep -i ubuntu | head -1 || echo "")
    if [[ -z "$UBUNTU_NAME" ]]; then
      log_info "安装 Ubuntu 24.04（请在弹出窗口设置用户名密码）..."
      wsl.exe --install -d Ubuntu-24.04 2>&1 || wsl.exe --install 2>&1 || {
        log_err "WSL2 安装失败，请手动执行: wsl --install -d Ubuntu-24.04"
        exit 1
      }
      UBUNTU_NAME="Ubuntu-24.04"
    else
      log "已有 WSL Ubuntu: $UBUNTU_NAME"
    fi

    # ─── 3d. 在 WSL2 内执行剩余安装 ───
    # 把整个安装脚本传到 WSL2 内部继续执行
    log_info "在 WSL2 Ubuntu 中继续安装..."

    # 将脚本自身复制到 WSL2
    SELF_SCRIPT="$KOGUN_HOME/install_wsl_inner.sh"
    cat > "$SELF_SCRIPT" << 'INNER'
#!/bin/bash
# === 在 WSL2 内部执行 ===
set -euo pipefail

# 颜色（复用）
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; NC='\033[0m'
log()     { echo -e "${GREEN}✅${NC} $1"; }
log_info(){ echo -e "${BLUE}ℹ️${NC} $1"; }
log_warn(){ echo -e "${YELLOW}⚠️${NC} $1"; }
log_err() { echo -e "${RED}❌${NC} $1"; }
step()    { echo -e "\n${CYAN}━━━ $1 ━━━${NC}"; }

KOGUN_HOME="$HOME/.kogun"
mkdir -p "$KOGUN_HOME"/{bin,logs,config,data}

# 1) 更新包管理器
log_info "更新 apt..."
apt-get update -qq 2>/dev/null || true
apt-get install -y -qq python3 python3-pip python3-venv curl wget jq 2>/dev/null || true

# 2) 安装 CUDA 工具
if ! command -v nvidia-smi &>/dev/null; then
  log_info "安装 WSL2 NVIDIA 驱动支持..."
  # WSL2 用 Windows 的 NVIDIA 驱动，只需安装 CUDA Toolkit
  wget -q https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb -O /tmp/cuda-keyring.deb 2>/dev/null || true
  if [[ -f /tmp/cuda-keyring.deb ]]; then
    dpkg -i /tmp/cuda-keyring.deb 2>/dev/null || true
    apt-get update -qq 2>/dev/null || true
    apt-get install -y -qq cuda-toolkit-12-6 2>/dev/null || apt-get install -y -qq nvidia-cuda-toolkit 2>/dev/null || true
    rm -f /tmp/cuda-keyring.deb
  fi
fi

if command -v nvidia-smi &>/dev/null; then
  GPU_INFO=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null | head -1)
  log "WSL2 GPU 已就绪: $GPU_INFO"
else
  log_warn "WSL2 中 GPU 不可见——确保 Windows 安装了 NVIDIA WSL2 驱动"
  log_info "下载: https://developer.nvidia.com/cuda/wsl"
fi

# 3) 安装 KOGUN CLI
log_info "安装 KOGUN CLI..."
CLI_URL="https://raw.githubusercontent.com/kogun/kogun-cli/main/kogun-cli.py"
if curl -sfL "$CLI_URL" -o "$KOGUN_HOME/bin/kogun" 2>/dev/null; then
  chmod +x "$KOGUN_HOME/bin/kogun"
  echo 'alias kogun="python3 $HOME/.kogun/bin/kogun"' >> "$HOME/.bashrc"
  log "CLI 已安装"
else
  log_err "CLI 下载失败"
fi

# 4) 安装 vLLM
if command -v nvidia-smi &>/dev/null; then
  log_info "安装 vLLM..."
  pip3 install --quiet --upgrade pip 2>/dev/null || true
  pip3 install --quiet vllm 2>/dev/null && log "vLLM 安装完成" || {
    pip3 install --quiet vllm==0.6.3 2>/dev/null && log "vLLM 0.6.3 安装完成" || log_warn "vLLM 安装失败"
  }
fi

# 5) 复制 vLLM 部署脚本
VLLM_SCRIPT_URL="https://raw.githubusercontent.com/kogun/vllm-deploy/main/vllm_deploy.sh"
curl -sfL "$VLLM_SCRIPT_URL" -o "$KOGUN_HOME/vllm_deploy.sh" 2>/dev/null || true
if [[ -f "$KOGUN_HOME/vllm_deploy.sh" ]]; then
  chmod +x "$KOGUN_HOME/vllm_deploy.sh"
  log "vLLM 部署脚本已就绪"
fi

# 完成
echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║${NC}  WSL2 环境安装完成！                           ${GREEN}║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════╝${NC}"
echo ""
echo "📟 在 WSL2 中执行:"
echo "  kogun register    注册设备"
echo "  kogun start       开始贡献算力"
echo "  bash ~/.kogun/vllm_deploy.sh --auto   部署推理服务"
echo ""
echo "💡 每次进入 WSL2: 打开 Windows 终端 → wsl"
INNER

    # 复制到 WSL2 并执行
    chmod +x "$SELF_SCRIPT"
    wsl.exe -d "$UBUNTU_NAME" bash -c "cat > /tmp/install_inner.sh" < "$SELF_SCRIPT" 2>/dev/null || {
      # 如果复制失败，通过 Windows 路径共享
      cp "$SELF_SCRIPT" /mnt/c/Users/Public/install_inner.sh 2>/dev/null || true
    }

    log_info "WSL2 安装脚本已就绪"
    echo ""
    echo -e "${CYAN}请在 WSL2 终端中执行以下命令完成安装：${NC}"
    echo ""
    if [[ -f "/mnt/c/Users/Public/install_inner.sh" ]]; then
      echo "  wsl"
      echo "  sudo bash /mnt/c/Users/Public/install_inner.sh"
    else
      echo "  wsl"
      echo "  bash /tmp/install_inner.sh"
    fi
    echo ""
    echo "或者直接进入 WSL2 执行:"
    echo "  wsl"
    echo "  sudo apt update && sudo apt install python3 python3-pip"
    echo "  pip install vllm"
    echo "  curl -sL https://kogun.ai/kogun-cli.py -o ~/.kogun/bin/kogun"
    echo ""

    exit 0
  else
    log_info "跳过 WSL2，仅安装 Windows CLI"
    log_info "Windows 只能贡献 CPU 算力"
  fi
fi

# ============================================================
# PHASE 3: 安装 CLI（Linux / macOS / WSL2 内部）
# ============================================================
if ! $ONLY_VLLM; then
  step "2/5 安装 KOGUN CLI"

  CLI_DEST="$KOGUN_HOME/bin/kogun"
  
  # 尝试多种方式获取 CLI
  CLI_FOUND=false
  for src in \
    "/opt/kogun/static/kogun-cli.py" \
    "/mnt/c/opt/kogun/static/kogun-cli.py" \
    "/opt/kogun/deploy/kogun-cli.py"; do
    if [[ -f "$src" ]]; then
      cp "$src" "$CLI_DEST"
      CLI_FOUND=true
      log "从 $src 复制 CLI"
      break
    fi
  done

  if ! $CLI_FOUND; then
    log_info "下载 KOGUN CLI..."
    curl -sL "https://raw.githubusercontent.com/kogun/kogun-cli/main/kogun-cli.py" -o "$CLI_DEST" 2>/dev/null || {
      log_err "下载失败"
      exit 1
    }
    log "CLI 下载完成"
  fi

  chmod +x "$CLI_DEST"
  echo "  CLI: $CLI_DEST"

  # 创建包装器
  cat > "$KOGUN_HOME/bin/kogun.sh" << 'WRAP'
#!/bin/bash
KOGUN_HOME="$HOME/.kogun"
for cmd in python3 python; do
  if command -v "$cmd" &>/dev/null && "$cmd" --version 2>&1 | grep -q "^Python 3"; then
    exec "$cmd" "$KOGUN_HOME/bin/kogun" "$@"
  fi
done
echo "❌ 需要 Python 3" >&2
exit 1
WRAP
  chmod +x "$KOGUN_HOME/bin/kogun.sh"

  # 配置 PATH
  SHELL_RC=""
  case "$SHELL" in
    */zsh) SHELL_RC="$HOME/.zshrc" ;;
    */bash) SHELL_RC="$HOME/.bashrc" ;;
  esac
  if [[ -n "$SHELL_RC" ]]; then
    if ! grep -q 'KOGUN' "$SHELL_RC" 2>/dev/null; then
      {
        echo ""
        echo "# KOGUN 贡献者节点"
        echo "export PATH=\"\$HOME/.kogun/bin:\$PATH\""
        echo "export KOGUN_API_URL=\"${KOGUN_API_URL}\""
        echo "alias kogun='\$HOME/.kogun/bin/kogun.sh'"
      } >> "$SHELL_RC"
      echo "  已配置 PATH 到 $SHELL_RC"
    fi
  fi

  # Windows 批处理
  if [[ "$OS" == "windows" ]]; then
    cat > "$KOGUN_HOME/bin/kogun.cmd" << 'BAT'
@echo off
set "KOGUN_HOME=%USERPROFILE%\.kogun"
python "%KOGUN_HOME%\bin\kogun" %*
BAT
    echo "  Windows 批处理: $KOGUN_HOME/bin/kogun.cmd"
  fi

  # 验证
  if $PYTHON "$CLI_DEST" version &>/dev/null; then
    log "CLI 验证通过: $($PYTHON "$CLI_DEST" version 2>&1 | head -1)"
  fi
fi

# ============================================================
# PHASE 4: 安装系统服务（Linux / macOS）
# ============================================================
if ! $ONLY_VLLM; then
  step "3/5 安装系统服务"

  if [[ "$OS" == "linux" ]] && command -v systemctl &>/dev/null; then
    # systemd 服务
    cat > "$KOGUN_HOME/config/kogun-contributor.service" << 'SRV'
[Unit]
Description=KOGUN 贡献者算力节点
After=network.target
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=%h/.kogun/bin/kogun.sh start
ExecStop=%h/.kogun/bin/kogun.sh stop
RemainAfterExit=yes
User=%i
Environment=KOGUN_API_URL=https://api.kogun.cn

[Install]
WantedBy=default.target
SRV
    echo "  🐧 systemd: $KOGUN_HOME/config/kogun-contributor.service"
    echo "  安装: sudo cp $KOGUN_HOME/config/kogun-contributor.service /etc/systemd/system/"
    echo "        sudo systemctl enable --now kogun-contributor"

  elif [[ "$OS" == "macos" ]]; then
    cat > "$KOGUN_HOME/config/com.kogun.contributor.plist" << 'PLIST'
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>Label</key>
    <string>com.kogun.contributor</string>
    <key>ProgramArguments</key>
    <array>
        <string>${HOME}/.kogun/bin/kogun.sh</string>
        <string>start</string>
    </array>
    <key>RunAtLoad</key>
    <true/>
    <key>KeepAlive</key>
    <false/>
    <key>EnvironmentVariables</key>
    <dict>
        <key>KOGUN_API_URL</key>
        <string>https://api.kogun.cn</string>
    </dict>
</dict>
</plist>
PLIST
    echo "  🍎 launchd: $KOGUN_HOME/config/com.kogun.contributor.plist"
  fi
fi

# ============================================================
# PHASE 5: 部署 vLLM（Linux + NVIDIA 自动执行，其他跳过）
# ============================================================
step "4/5 部署推理服务"

if $HAS_NVIDIA && [[ "$OS" == "linux" ]]; then
  echo -e "  ${GREEN}检测到 NVIDIA GPU: $GPU_NAME ($GPU_VRAM_MB MB)${NC}"
  echo ""

  # 找 vLLM 部署脚本
  VLLM_SCRIPT=""
  for p in \
    "/opt/kogun/deploy/vllm_deploy.sh" \
    "/mnt/c/opt/kogun/deploy/vllm_deploy.sh" \
    "$KOGUN_HOME/vllm_deploy.sh"; do
    if [[ -f "$p" ]]; then
      VLLM_SCRIPT="$p"
      break
    fi
  done

  if [[ -n "$VLLM_SCRIPT" ]]; then
    if $AUTO_MODE; then
      log_info "自动部署 vLLM..."
      bash "$VLLM_SCRIPT" --auto 2>&1 | tail -10 || log_warn "vLLM 自动部署失败"
    else
      echo -e "${CYAN}你的 GPU 可以跑 vLLM 推理服务！${NC}"
      echo "  按显存推荐:"
      if [[ "$GPU_VRAM_MB" -ge 8192 ]]; then
        echo "    • Qwen2.5-7B（4GB）  ·  Llama-3.1-8B（6GB）  ·  DeepSeek-Coder-V2-Lite（8GB）"
      elif [[ "$GPU_VRAM_MB" -ge 4096 ]]; then
        echo "    • Qwen2.5-7B（4GB）"
      else
        echo "    ${YELLOW}显存不足 4GB，无法部署 vLLM${NC}"
        VLLM_SCRIPT=""
      fi
      echo ""
      read -rp "立即部署 vLLM? [Y/n]: " DO_VLLM
      if [[ ! "$DO_VLLM" =~ ^[Nn] ]]; then
        log_info "部署 vLLM..."
        bash "$VLLM_SCRIPT" --auto 2>&1 | tail -15 || log_warn "自动部署遇到问题，手动执行: bash $VLLM_SCRIPT --auto"
      fi
    fi
  else
    if $AUTO_MODE; then
      log_info "vLLM 脚本未找到，跳过"
    else
      echo "  一键部署脚本: bash vllm_deploy.sh --auto"
      echo "  （需从平台下载）"
    fi
  fi
elif [[ "$OS" == "linux" ]]; then
  echo "  ${YELLOW}无 NVIDIA GPU——跳过 vLLM 部署${NC}"
  echo "  仍可贡献 CPU 算力"
elif [[ "$OS" == "macos" ]]; then
  echo "  ${YELLOW}macOS 不支持 vLLM${NC}"
  echo "  仍可贡献 CPU 算力"
fi

# ============================================================
# PHASE 6: 设备注册
# ============================================================
if ! $ONLY_VLLM; then
  step "5/5 设备注册"

  if $AUTO_MODE; then
    log_info "自动模式已开启，跳过注册"
    echo "  安装后手动注册: kogun register"
  else
    echo -e "${CYAN}要现在注册设备到 KOGUN 网络吗？${NC}"
    echo "  注册后可一键启停、查看收益"
    echo ""
    read -rp "立即注册? [Y/n]: " DO_REG
    if [[ ! "$DO_REG" =~ ^[Nn] ]]; then
      echo ""
      $PYTHON "$CLI_DEST" register
    fi
  fi
fi

# ============================================================
# 完成
# ============================================================
echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║${NC}  ${CYAN}KOGUN 贡献者节点安装完成！${NC}                           ${GREEN}║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════════╝${NC}"
echo ""
echo "  📟 命令:   kogun <register|start|stop|status|earnings|vllm>"
echo "  📂 CLI:    $KOGUN_HOME/bin/kogun"
echo "  📂 配置:   $KOGUN_HOME/config/"
echo "  📂 日志:   $KOGUN_HOME/logs/"
echo "  🌐 API:    $KOGUN_API_URL"
echo ""

if [[ -n "$SHELL_RC" ]]; then
  echo "  ⚡ 执行: source $SHELL_RC   （然后可直接运行 kogun）"
fi
echo ""
echo "📟 常用命令:"
echo "  kogun register    注册设备"
echo "  kogun status      查看设备状态 + 收益"
echo "  kogun start       开始贡献算力"
echo "  kogun earnings    查看收益明细"
echo ""

if $HAS_NVIDIA && [[ "$OS" == "linux" ]] && ! $ONLY_VLLM; then
  echo -e "⚡ ${CYAN}你的 GPU 已就绪！${NC}"
  echo "  部署 vLLM 推理服务:"
  echo "    bash install.sh --vllm"
  echo ""
fi