#!/usr/bin/env bash
# Host dependency setup for the DgxManager platform .deb.
# This script is intentionally best-effort: DgxManager itself may install even
# when Docker/NVIDIA runtime setup needs operator follow-up.
set -u

LOG_FILE="${DGXMANAGER_HOST_SETUP_LOG:-/var/log/dgxmanager/install-host-setup.log}"
ASSUME_YES="${DGXMANAGER_HOST_SETUP_ASSUME_YES:-1}"
SKIP="${DGXMANAGER_SKIP_HOST_SETUP:-0}"
DOCKER_APT_MIRRORS="${DGXMANAGER_DOCKER_APT_MIRRORS:-https://mirrors.tuna.tsinghua.edu.cn/docker-ce,https://mirrors.ustc.edu.cn/docker-ce,https://mirrors.aliyun.com/docker-ce}"
NVIDIA_APT_MIRRORS="${DGXMANAGER_NVIDIA_APT_MIRRORS:-https://mirrors.tuna.tsinghua.edu.cn/nvidia-container-toolkit,https://mirrors.ustc.edu.cn/nvidia-container-toolkit}"
DOCKER_REGISTRY_MIRRORS="${DGXMANAGER_DOCKER_REGISTRY_MIRRORS:-https://docker.1ms.run,https://docker.xuanyuan.me,https://docker.1panel.live,https://docker.1panel.dev,https://docker.m.daocloud.io,https://dockerpull.com,https://dockerhub.icu,https://hub-mirror.c.163.com,https://mirror.baidubce.com,https://docker.mirrors.ustc.edu.cn}"
export DGXMANAGER_DOCKER_REGISTRY_MIRRORS="$DOCKER_REGISTRY_MIRRORS"
WARNINGS=""
DOCKER_RESTART_NEEDED=0
DOCKER_GROUP_CHANGED=0

mkdir -p "$(dirname "$LOG_FILE")"
: > "$LOG_FILE"

log() {
  printf '%s %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*" | tee -a "$LOG_FILE"
}

warn() {
  WARNINGS="${WARNINGS}\n- $*"
  log "WARNING: $*"
}

run() {
  log "+ $*"
  "$@" >>"$LOG_FILE" 2>&1
}

have_cmd() {
  command -v "$1" >/dev/null 2>&1
}

is_deb_installed() {
  dpkg-query -W -f='${Status}' "$1" 2>/dev/null | grep -q 'install ok installed'
}

apt_update() {
  if [ "${APT_UPDATED:-0}" = "1" ]; then
    return 0
  fi
  if run apt-get update; then
    APT_UPDATED=1
    return 0
  fi
  return 1
}

apt_install() {
  apt_update || return 1
  DEBIAN_FRONTEND=noninteractive apt-get install -y "$@" >>"$LOG_FILE" 2>&1
}

ubuntu_codename() {
  if [ -r /etc/os-release ]; then
    . /etc/os-release
    if [ -n "${VERSION_CODENAME:-}" ]; then
      printf '%s\n' "$VERSION_CODENAME"
      return 0
    fi
    case "${VERSION_ID:-}" in
      24.04) printf 'noble\n'; return 0 ;;
      22.04) printf 'jammy\n'; return 0 ;;
      20.04) printf 'focal\n'; return 0 ;;
    esac
  fi
  return 1
}

normalize_arch() {
  case "$(dpkg --print-architecture 2>/dev/null || uname -m)" in
    amd64|x86_64) printf 'amd64\n' ;;
    arm64|aarch64) printf 'arm64\n' ;;
    *) dpkg --print-architecture 2>/dev/null || uname -m ;;
  esac
}

fetch_url_to_file() {
  url="$1"
  out="$2"
  if have_cmd curl; then
    curl -fsSL "$url" -o "$out" >>"$LOG_FILE" 2>&1
    return $?
  fi
  if have_cmd wget; then
    wget -qO "$out" "$url" >>"$LOG_FILE" 2>&1
    return $?
  fi
  return 127
}

install_download_tools() {
  have_cmd curl && have_cmd gpg && return 0
  apt_install ca-certificates curl gnupg || return 1
}

configure_docker_ce_cn_repo() {
  codename="$(ubuntu_codename 2>/dev/null || true)"
  arch="$(normalize_arch)"
  if [ -z "$codename" ]; then
    warn "无法识别 Ubuntu 代号，跳过 Docker CE 国内 apt 源配置。"
    return 1
  fi
  install_download_tools || {
    warn "无法安装 curl/gnupg，跳过 Docker CE 国内 apt 源配置。"
    return 1
  }
  mkdir -p /etc/apt/keyrings
  for mirror in $(printf '%s' "$DOCKER_APT_MIRRORS" | tr ',' ' '); do
    mirror="$(printf '%s' "$mirror" | sed 's:/*$::')"
    [ -n "$mirror" ] || continue
    log "尝试配置 Docker CE 国内源: $mirror"
    if fetch_url_to_file "$mirror/linux/ubuntu/gpg" /tmp/dgxmanager-docker.gpg; then
      if gpg --dearmor < /tmp/dgxmanager-docker.gpg > /etc/apt/keyrings/dgxmanager-docker.gpg 2>>"$LOG_FILE"; then
        printf 'deb [arch=%s signed-by=/etc/apt/keyrings/dgxmanager-docker.gpg] %s/linux/ubuntu %s stable\n' "$arch" "$mirror" "$codename" > /etc/apt/sources.list.d/dgxmanager-docker.list
        APT_UPDATED=0
        if apt_update; then
          log "Docker CE 国内源配置完成: $mirror"
          return 0
        fi
      fi
    fi
    warn "Docker CE 国内源不可用: $mirror"
  done
  return 1
}

configure_nvidia_toolkit_cn_repo() {
  install_download_tools || {
    warn "无法安装 curl/gnupg，跳过 NVIDIA Container Toolkit 国内 apt 源配置。"
    return 1
  }
  mkdir -p /etc/apt/keyrings
  for mirror in $(printf '%s' "$NVIDIA_APT_MIRRORS" | tr ',' ' '); do
    mirror="$(printf '%s' "$mirror" | sed 's:/*$::')"
    [ -n "$mirror" ] || continue
    log "尝试配置 NVIDIA Container Toolkit 国内源: $mirror"
    if fetch_url_to_file "$mirror/gpgkey" /tmp/dgxmanager-nvidia-container-toolkit.gpg; then
      if gpg --dearmor < /tmp/dgxmanager-nvidia-container-toolkit.gpg > /etc/apt/keyrings/dgxmanager-nvidia-container-toolkit.gpg 2>>"$LOG_FILE"; then
        printf 'deb [signed-by=/etc/apt/keyrings/dgxmanager-nvidia-container-toolkit.gpg] %s/stable/deb/$(ARCH) /\n' "$mirror" > /etc/apt/sources.list.d/dgxmanager-nvidia-container-toolkit.list
        APT_UPDATED=0
        if apt_update; then
          log "NVIDIA Container Toolkit 国内源配置完成: $mirror"
          return 0
        fi
      fi
    fi
    warn "NVIDIA Container Toolkit 国内源不可用: $mirror"
  done
  return 1
}

install_docker_stack() {
  if have_cmd docker; then
    log "Docker CLI 已存在: $(docker --version 2>/dev/null || true)"
  else
    log "未检测到 Docker，优先从当前 apt 源安装 docker.io。"
    if apt_install docker.io; then
      log "docker.io 安装完成。"
    else
      warn "当前 apt 源安装 docker.io 失败，尝试 Docker CE 国内镜像源。"
      if configure_docker_ce_cn_repo && apt_install docker-ce docker-ce-cli containerd.io; then
        log "Docker CE 安装完成。"
      else
        warn "Docker 安装失败。请确认 apt 源和网络后重装，或手动安装 Docker。"
      fi
    fi
  fi

  if docker compose version >/dev/null 2>&1; then
    log "Docker Compose plugin 已可用: $(docker compose version 2>/dev/null || true)"
  else
    if apt_install docker-compose-plugin; then
      log "docker-compose-plugin 安装完成。"
    else
      warn "docker-compose-plugin 安装失败；部署容器运行时可能不可用。"
    fi
  fi
}

configure_docker_daemon_mirrors() {
  mkdir -p /etc/docker
  python3 - <<'PY' >>"$LOG_FILE" 2>&1
from __future__ import annotations
import json
import os
from pathlib import Path
path = Path('/etc/docker/daemon.json')
mirrors = []
seen = set()
for raw in os.environ.get('DGXMANAGER_DOCKER_REGISTRY_MIRRORS', '').split(','):
    value = raw.strip().rstrip('/')
    if value and value.startswith(('http://', 'https://')) and value not in seen:
        mirrors.append(value)
        seen.add(value)
if not mirrors:
    mirrors = [
        'https://docker.1ms.run',
        'https://docker.xuanyuan.me',
        'https://docker.1panel.live',
        'https://docker.1panel.dev',
        'https://docker.m.daocloud.io',
        'https://dockerpull.com',
        'https://dockerhub.icu',
        'https://hub-mirror.c.163.com',
        'https://mirror.baidubce.com',
        'https://docker.mirrors.ustc.edu.cn',
    ]
try:
    data = json.loads(path.read_text(encoding='utf-8')) if path.exists() and path.read_text(encoding='utf-8').strip() else {}
except Exception:
    backup = path.with_suffix('.json.dgxmanager-bak')
    if path.exists():
        backup.write_text(path.read_text(encoding='utf-8', errors='ignore'), encoding='utf-8')
    data = {}
if not isinstance(data, dict):
    data = {}
old = data.get('registry-mirrors')
data['registry-mirrors'] = mirrors
tmp = path.with_suffix('.json.tmp')
tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + '\n', encoding='utf-8')
tmp.replace(path)
print('docker daemon registry mirrors:', ', '.join(mirrors))
print('previous registry mirrors:', old)
PY
  if [ $? -eq 0 ]; then
    log "Docker registry mirror 已写入 /etc/docker/daemon.json。"
    DOCKER_RESTART_NEEDED=1
  else
    warn "写入 Docker registry mirror 失败。"
  fi
}

start_docker_service() {
  if ! have_cmd systemctl; then
    warn "未检测到 systemctl，跳过 Docker 服务启动。"
    return 0
  fi
  if systemctl list-unit-files docker.service >/dev/null 2>&1; then
    run systemctl enable docker.service || warn "启用 docker.service 失败。"
    run systemctl start docker.service || warn "启动 docker.service 失败。"
  fi
}

restart_docker_if_needed() {
  [ "$DOCKER_RESTART_NEEDED" = "1" ] || return 0
  if have_cmd systemctl && systemctl list-unit-files docker.service >/dev/null 2>&1; then
    run systemctl restart docker.service || warn "重启 docker.service 失败，请手动执行 systemctl restart docker。"
  else
    warn "Docker 配置已变更，但无法自动重启 Docker。"
  fi
}

resolve_target_user() {
  if [ -n "${DGXMANAGER_INSTALL_USER:-}" ] && id "$DGXMANAGER_INSTALL_USER" >/dev/null 2>&1; then
    printf '%s\n' "$DGXMANAGER_INSTALL_USER"
    return 0
  fi
  if [ -n "${SUDO_USER:-}" ] && [ "$SUDO_USER" != "root" ] && id "$SUDO_USER" >/dev/null 2>&1; then
    printf '%s\n' "$SUDO_USER"
    return 0
  fi
  if [ -n "${PKEXEC_UID:-}" ]; then
    user="$(getent passwd "$PKEXEC_UID" | cut -d: -f1 || true)"
    if [ -n "$user" ] && [ "$user" != "root" ]; then
      printf '%s\n' "$user"
      return 0
    fi
  fi
  if have_cmd loginctl; then
    user="$(loginctl list-sessions --no-legend 2>/dev/null | awk '$4 ~ /seat/ && $3 != "root" {print $3; exit}')"
    if [ -n "$user" ] && id "$user" >/dev/null 2>&1; then
      printf '%s\n' "$user"
      return 0
    fi
  fi
  user="$(find /home -mindepth 1 -maxdepth 1 -type d -printf '%f\n' 2>/dev/null | head -n1)"
  if [ -n "$user" ] && id "$user" >/dev/null 2>&1; then
    printf '%s\n' "$user"
    return 0
  fi
  return 1
}

configure_docker_group() {
  if ! getent group docker >/dev/null 2>&1; then
    run groupadd docker || warn "创建 docker 用户组失败。"
  fi
  user="$(resolve_target_user 2>/dev/null || true)"
  if [ -z "$user" ]; then
    warn "无法识别桌面用户，跳过加入 docker 组。可手动执行: sudo usermod -aG docker <用户名>。"
    return 0
  fi
  if id -nG "$user" 2>/dev/null | tr ' ' '\n' | grep -qx docker; then
    log "用户 $user 已在 docker 组。"
  else
    if run usermod -aG docker "$user"; then
      DOCKER_GROUP_CHANGED=1
      log "已将用户 $user 加入 docker 组。该权限通常需要注销重登后生效。"
    else
      warn "将用户 $user 加入 docker 组失败。"
    fi
  fi
}

install_nvidia_container_toolkit() {
  if have_cmd nvidia-ctk && is_deb_installed nvidia-container-toolkit; then
    log "nvidia-container-toolkit 已存在: $(nvidia-ctk --version 2>/dev/null || true)"
    return 0
  fi
  log "尝试从当前 apt 源安装 nvidia-container-toolkit。"
  if apt_install nvidia-container-toolkit; then
    log "nvidia-container-toolkit 安装完成。"
    return 0
  fi
  warn "当前 apt 源安装 nvidia-container-toolkit 失败，尝试国内镜像源。"
  if configure_nvidia_toolkit_cn_repo && apt_install nvidia-container-toolkit; then
    log "nvidia-container-toolkit 通过国内源安装完成。"
    return 0
  fi
  warn "nvidia-container-toolkit 安装失败。GPU 容器部署需要安装后再运行 nvidia-ctk runtime configure --runtime=docker。"
  return 1
}

configure_nvidia_runtime() {
  if ! have_cmd nvidia-smi; then
    warn "未检测到 nvidia-smi。DgxManager 已安装，但 GPU 推理需要先安装 NVIDIA 驱动并重启系统。"
    return 0
  fi
  log "NVIDIA 驱动检测: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1 || true)"
  if ! have_cmd nvidia-ctk; then
    install_nvidia_container_toolkit || return 0
  fi
  if have_cmd nvidia-ctk; then
    if run nvidia-ctk runtime configure --runtime=docker; then
      log "Docker NVIDIA runtime 注册完成。"
      DOCKER_RESTART_NEEDED=1
    else
      warn "nvidia-ctk runtime configure 执行失败。"
    fi
  fi
}

verify_summary() {
  if have_cmd docker; then
    if docker info >/tmp/dgxmanager-docker-info.txt 2>&1; then
      log "Docker daemon 可用。"
      if docker info --format '{{json .Runtimes}}' 2>/dev/null | grep -qi nvidia; then
        log "Docker nvidia runtime 可用。"
      elif have_cmd nvidia-smi; then
        warn "Docker 尚未报告 nvidia runtime；GPU 容器部署可能不可用。"
      fi
    else
      warn "Docker CLI 存在，但 daemon 当前不可用。详情见 $LOG_FILE。"
      cat /tmp/dgxmanager-docker-info.txt >>"$LOG_FILE" 2>/dev/null || true
    fi
  else
    warn "Docker CLI 不可用。"
  fi
  if [ "$DOCKER_GROUP_CHANGED" = "1" ]; then
    warn "已更新 docker 用户组。请注销并重新登录，使当前用户的非 root Docker 权限生效。"
  fi
  if [ -n "$WARNINGS" ]; then
    printf '\nDgxManager host setup completed with notices:%b\n' "$WARNINGS"
    printf 'Detailed log: %s\n' "$LOG_FILE"
  else
    printf '\nDgxManager host setup completed successfully. Detailed log: %s\n' "$LOG_FILE"
  fi
}

main() {
  if [ "$SKIP" = "1" ] || [ "$SKIP" = "true" ]; then
    log "DGXMANAGER_SKIP_HOST_SETUP=1，跳过 Docker/NVIDIA 主机环境配置。"
    return 0
  fi
  if [ "$(id -u)" != "0" ]; then
    warn "主机环境配置需要 root；当前不是 root，跳过。"
    verify_summary
    return 0
  fi
  if ! have_cmd apt-get || ! have_cmd dpkg-query; then
    warn "当前系统不是 apt/dpkg 环境，跳过自动安装 Docker/NVIDIA 依赖。"
    verify_summary
    return 0
  fi
  log "开始配置 DgxManager 主机环境（Docker / NVIDIA Container Toolkit / 国内镜像）。"
  install_docker_stack
  configure_docker_daemon_mirrors
  start_docker_service
  configure_docker_group
  configure_nvidia_runtime
  restart_docker_if_needed
  verify_summary
  return 0
}

main "$@"
exit 0
