#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RUN_DIR="$(pwd)"

ROCM_TEMPLATE="${SCRIPT_DIR}/templates/hip-libraries-rocm-ubuntu-630.Dockerfile"
CUDA_TEMPLATE="${SCRIPT_DIR}/templates/hip-libraries-cuda-ubuntu-630.Dockerfile"

ROCM_DOCKERFILE="${RUN_DIR}/hip-libraries-rocm-ubuntu-630.Dockerfile"
CUDA_DOCKERFILE="${RUN_DIR}/hip-libraries-cuda-ubuntu-630.Dockerfile"

ROCM_WHEEL="${SCRIPT_DIR}/nwavesdk-0.0.2a0+rocm-cp310-cp310-linux_x86_64.whl"
CUDA_WHEEL="${SCRIPT_DIR}/nwavesdk-0.0.2a0+cu-cp310-cp310-linux_x86_64.whl"

usage() {
  cat <<'EOF_USAGE'
Usage:
  ./install.sh \
    --vendor amd|nvidia \
    --dockermem-dir /absolute/path \
    --image-name IMAGE_NAME \
    --container-name CONTAINER_NAME \
    [--amd-gfx 10.3.0|11.0.0]

Required:
  --vendor           GPU vendor: amd or nvidia
  --dockermem-dir    Absolute host path where DockerMem workspace root lives (host path where project files using nwave must be placed)
  --image-name       Docker image tag to build
  --container-name   Docker container name to run

Optional:
  --amd-gfx          AMD HSA override. Default: 11.0.0
                     Typical values:
                       10.3.0 -> RDNA / RDNA2
                       11.0.0 -> RDNA3 / RDNA3.5

Notes:
  The script expects these wheels in the same directory as install.sh:
    - nwavesdk-0.0.2a0+cu-cp310-cp310-linux_x86_64.whl
    - nwavesdk-0.0.2a0+rocm-cp310-cp310-linux_x86_64.whl

Examples:
  ./install.sh \
    --vendor amd \
    --dockermem-dir /home/giuseppe/DockerMem \
    --image-name kernel_devel_amd_630 \
    --container-name nwave_amd_dev \
    --amd-gfx 11.0.0

  ./install.sh \
    --vendor nvidia \
    --dockermem-dir /home/giuseppe/DockerMem \
    --image-name kernel_devel_nvidia_630 \
    --container-name nwave_nvidia_dev
EOF_USAGE
}

VENDOR=""
DOCKERMEM_DIR=""
IMAGE_NAME=""
CONTAINER_NAME=""
AMD_GFX="11.0.0"
SELECTED_WHEEL=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --vendor)
      VENDOR="$2"
      shift 2
      ;;
    --dockermem-dir)
      DOCKERMEM_DIR="$2"
      shift 2
      ;;
    --image-name)
      IMAGE_NAME="$2"
      shift 2
      ;;
    --container-name)
      CONTAINER_NAME="$2"
      shift 2
      ;;
    --amd-gfx)
      AMD_GFX="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Error: unknown argument: $1" >&2
      usage
      exit 1
      ;;
  esac
done

if [[ -z "${VENDOR}" || -z "${DOCKERMEM_DIR}" || -z "${IMAGE_NAME}" || -z "${CONTAINER_NAME}" ]]; then
  echo "Error: missing required arguments." >&2
  usage
  exit 1
fi

if [[ "${VENDOR}" != "amd" && "${VENDOR}" != "nvidia" ]]; then
  echo "Error: --vendor must be 'amd' or 'nvidia'." >&2
  exit 1
fi

if [[ "${DOCKERMEM_DIR}" != /* ]]; then
  echo "Error: --dockermem-dir must be an absolute path." >&2
  exit 1
fi

if ! command -v docker >/dev/null 2>&1; then
  echo "Error: docker is not installed or not in PATH." >&2
  exit 1
fi

HOST_UID="$(id -u)"
RENDER_GID="$(getent group render | cut -d: -f3 || true)"

if [[ -z "${RENDER_GID}" ]]; then
  echo "Error: could not determine render group GID. Is the 'render' group present on the host?" >&2
  exit 1
fi

WORKSPACE_DIR="${DOCKERMEM_DIR%/}/${CONTAINER_NAME}"

mkdir -p "${WORKSPACE_DIR}"

select_wheel() {
  if [[ "${VENDOR}" == "amd" ]]; then
    SELECTED_WHEEL="${ROCM_WHEEL}"
  else
    SELECTED_WHEEL="${CUDA_WHEEL}"
  fi

  if [[ ! -f "${SELECTED_WHEEL}" ]]; then
    echo "Error: missing wheel for vendor '${VENDOR}': ${SELECTED_WHEEL}" >&2
    exit 1
  fi
}

echo "[INFO] Host UID       : ${HOST_UID}"
echo "[INFO] Render GID     : ${RENDER_GID}"
echo "[INFO] Workspace dir  : ${WORKSPACE_DIR}"
echo "[INFO] Run directory  : ${RUN_DIR}"

copy_templates() {
  if [[ ! -f "${ROCM_TEMPLATE}" ]]; then
    echo "Error: missing template: ${ROCM_TEMPLATE}" >&2
    exit 1
  fi

  if [[ ! -f "${CUDA_TEMPLATE}" ]]; then
    echo "Error: missing template: ${CUDA_TEMPLATE}" >&2
    exit 1
  fi

  cp "${ROCM_TEMPLATE}" "${ROCM_DOCKERFILE}"
  cp "${CUDA_TEMPLATE}" "${CUDA_DOCKERFILE}"

  echo "[INFO] Dockerfiles copied to current directory:"
  echo "       ${ROCM_DOCKERFILE}"
  echo "       ${CUDA_DOCKERFILE}"
}

patch_common_user_config() {
  local file="$1"

  sed -i -E "s/^ARG GID=.*/ARG GID=${RENDER_GID}/" "${file}"
  sed -i -E "s/^ARG UID=.*/ARG UID=${HOST_UID}/" "${file}"
}

patch_amd_gfx() {
  local file="$1"
  local gfx="$2"

  if grep -q '^ENV HSA_OVERRIDE_GFX_VERSION=' "${file}"; then
    sed -i -E "s|^ENV HSA_OVERRIDE_GFX_VERSION=.*|ENV HSA_OVERRIDE_GFX_VERSION=${gfx}|" "${file}"
  else
    printf '\nENV HSA_OVERRIDE_GFX_VERSION=%s\n' "${gfx}" >> "${file}"
  fi
}

build_image() {
  local dockerfile="$1"
  echo "[INFO] Building image '${IMAGE_NAME}' using ${dockerfile}"
  docker build "${RUN_DIR}" -f "${dockerfile}" -t "${IMAGE_NAME}"
}

remove_existing_container_if_any() {
  if docker ps -a --format '{{.Names}}' | grep -Fxq "${CONTAINER_NAME}"; then
    echo "[INFO] Container '${CONTAINER_NAME}' already exists. Removing it."
    docker rm -f "${CONTAINER_NAME}" >/dev/null
  fi
}

container_start_command() {
  local wheel_basename
  wheel_basename="$(basename "${SELECTED_WHEEL}")"

  printf '%s' "source /opt/conda/etc/profile.d/conda.sh && conda activate PyTorch && pip install --force-reinstall /tmp/nwave-installer/${wheel_basename} && exec bash -l"
}

run_amd() {
  remove_existing_container_if_any
  echo "[INFO] Running AMD container '${CONTAINER_NAME}'"
  echo "[INFO] Installing wheel: $(basename "${SELECTED_WHEEL}")"
  docker run \
    --mount type=bind,source="${WORKSPACE_DIR}",target=/home/workspaces \
    --mount type=bind,source="${SCRIPT_DIR}",target=/tmp/nwave-installer,readonly \
    -ti \
    --cap-add=SYS_PTRACE \
    --security-opt seccomp=unconfined \
    --device=/dev/kfd \
    --device=/dev/dri \
    --group-add video \
    --ipc=host \
    --shm-size 16G \
    --name "${CONTAINER_NAME}" \
    "${IMAGE_NAME}" \
    bash -lc "$(container_start_command)"
}

run_nvidia() {
  remove_existing_container_if_any
  echo "[INFO] Running NVIDIA container '${CONTAINER_NAME}'"
  echo "[INFO] Installing wheel: $(basename "${SELECTED_WHEEL}")"
  docker run \
    --mount type=bind,source="${WORKSPACE_DIR}",target=/home/workspaces \
    --mount type=bind,source="${SCRIPT_DIR}",target=/tmp/nwave-installer,readonly \
    -ti \
    --gpus all \
    --shm-size 16G \
    --name "${CONTAINER_NAME}" \
    "${IMAGE_NAME}" \
    bash -lc "$(container_start_command)"
}

select_wheel
copy_templates

patch_common_user_config "${ROCM_DOCKERFILE}"
patch_common_user_config "${CUDA_DOCKERFILE}"

if [[ "${VENDOR}" == "amd" ]]; then
  patch_amd_gfx "${ROCM_DOCKERFILE}" "${AMD_GFX}"
  build_image "${ROCM_DOCKERFILE}"
  run_amd
else
  build_image "${CUDA_DOCKERFILE}"
  run_nvidia
fi