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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [[ -d "${SCRIPT_DIR}/src/robot_arm_hardware" ]]; then
  DEFAULT_WORKSPACE="${SCRIPT_DIR}"
else
  DEFAULT_WORKSPACE="${HOME}/zeroarm_ros2_sdk"
fi

WORKSPACE="${ZEROARM_ROS_WORKSPACE:-${DEFAULT_WORKSPACE}}"
ROS_DISTRO="${ZEROARM_ROS_DISTRO:-humble}"
CPP_SDK_PREFIX="${ZEROARM_CPP_SDK_PREFIX:-/usr/local}"
CPP_SDK_BASE_URL="${CPP_SDK_BASE_URL:-https://robotic.forceease.tech/cpp-sdk/current}"
ROS_SDK_BASE_URL="${ROS_SDK_BASE_URL:-https://robotic.forceease.tech/ros-sdk/current}"
ROS_SDK_ARCHIVE="${ZEROARM_ROS_SDK_ARCHIVE:-zeroArm_ros2_sdk.tar.gz}"
BUILD_TYPE="${ZEROARM_ROS_BUILD_TYPE:-RelWithDebInfo}"
SKIP_CPP_SDK=0
SKIP_ROSDEP="${ZEROARM_SKIP_ROSDEP:-0}"
DRY_RUN=0
TMP_DIR=""

PACKAGES=(
  robot_arm_interfaces
  robot_arm_hardware
  robot_arm_description
  robot_arm_moveit
)

usage() {
  cat <<'USAGE'
Usage:
  bash scripts/install_zeroarm_ros2_sdk.sh [options]

Options:
  --workspace <path>   ROS2 workspace path. Defaults to the bundled workspace,
                       or ~/zeroarm_ros2_sdk for the standalone website script.
  --prefix <path>      C++ SDK install prefix. Defaults to /usr/local.
  --ros-distro <name>  ROS distro to source. Defaults to humble.
  --build-type <type>  CMake build type. Defaults to RelWithDebInfo.
  --skip-cpp-sdk       Do not download or install the C++ SDK binary.
  --skip-rosdep        Do not run rosdep install.
  --dry-run            Print planned commands without changing the system.
  -h, --help           Show this help.

Environment:
  ZEROARM_ROS_WORKSPACE
  ZEROARM_CPP_SDK_PREFIX
  ZEROARM_ROS_DISTRO
  ZEROARM_ROS_BUILD_TYPE
  ZEROARM_SKIP_ROSDEP
  ZEROARM_SKIP_DEPS     Passed to the C++ SDK installer.
  CPP_SDK_BASE_URL      Defaults to https://robotic.forceease.tech/cpp-sdk/current.
  ROS_SDK_BASE_URL      Defaults to https://robotic.forceease.tech/ros-sdk/current.
  ZEROARM_ROS_SDK_ARCHIVE
USAGE
}

log() {
  printf '[zeroarm-ros2-sdk] %s\n' "$*"
}

die() {
  printf '[zeroarm-ros2-sdk] ERROR: %s\n' "$*" >&2
  exit 1
}

run() {
  if [[ "${DRY_RUN}" == "1" ]]; then
    printf '+'
    printf ' %q' "$@"
    printf '\n'
    return 0
  fi
  "$@"
}

cleanup() {
  if [[ -n "${TMP_DIR}" && -d "${TMP_DIR}" ]]; then
    rm -rf "${TMP_DIR}"
  fi
}
trap cleanup EXIT

resolve_dir() {
  local path="$1"
  [[ -d "${path}" ]] || die "directory does not exist: ${path}"
  (cd -- "${path}" && pwd)
}

workspace_has_sources() {
  [[ -d "${WORKSPACE}/src/robot_arm_interfaces" ]] &&
    [[ -d "${WORKSPACE}/src/robot_arm_hardware" ]] &&
    [[ -d "${WORKSPACE}/src/robot_arm_description" ]] &&
    [[ -d "${WORKSPACE}/src/robot_arm_moveit" ]]
}

bootstrap_workspace() {
  if workspace_has_sources; then
    return 0
  fi

  local archive_url checksum_url archive_path checksum_path staging source_root
  archive_url="${ROS_SDK_BASE_URL%/}/${ROS_SDK_ARCHIVE}"
  checksum_url="${archive_url}.sha256"
  log "Bootstrapping ROS2 SDK workspace from ${archive_url}"
  TMP_DIR="$(mktemp -d)"
  archive_path="${TMP_DIR}/${ROS_SDK_ARCHIVE}"
  checksum_path="${archive_path}.sha256"
  staging="${TMP_DIR}/workspace"

  run mkdir -p "${WORKSPACE}" "${staging}"
  run curl -fL --retry 3 --connect-timeout 10 -o "${archive_path}" "${archive_url}"
  run curl -fL --retry 3 --connect-timeout 10 -o "${checksum_path}" "${checksum_url}"
  if [[ "${DRY_RUN}" == "1" ]]; then
    run sha256sum -c "${checksum_path}"
    run tar -xzf "${archive_path}" -C "${staging}"
    log "Would populate ${WORKSPACE} from the verified ROS2 SDK archive"
    return 0
  fi

  (cd -- "${TMP_DIR}" && sha256sum -c "$(basename -- "${checksum_path}")") ||
    die "ROS2 SDK archive checksum verification failed"
  tar -xzf "${archive_path}" -C "${staging}"
  if [[ -d "${staging}/zeroArm_ros2_sdk/src" ]]; then
    source_root="${staging}/zeroArm_ros2_sdk"
  elif [[ -d "${staging}/src" ]]; then
    source_root="${staging}"
  else
    die "ROS2 SDK archive does not contain a src workspace"
  fi
  cp -a "${source_root}/." "${WORKSPACE}/"
  workspace_has_sources || die "downloaded ROS2 SDK workspace is incomplete"
}

parse_args() {
  while [[ "$#" -gt 0 ]]; do
    case "$1" in
      --workspace)
        [[ "$#" -ge 2 ]] || die "--workspace requires a path"
        WORKSPACE="$2"
        shift 2
        ;;
      --prefix)
        [[ "$#" -ge 2 ]] || die "--prefix requires a path"
        CPP_SDK_PREFIX="$2"
        shift 2
        ;;
      --ros-distro)
        [[ "$#" -ge 2 ]] || die "--ros-distro requires a value"
        ROS_DISTRO="$2"
        shift 2
        ;;
      --build-type)
        [[ "$#" -ge 2 ]] || die "--build-type requires a value"
        BUILD_TYPE="$2"
        shift 2
        ;;
      --skip-cpp-sdk)
        SKIP_CPP_SDK=1
        shift
        ;;
      --skip-rosdep)
        SKIP_ROSDEP=1
        shift
        ;;
      --dry-run)
        DRY_RUN=1
        shift
        ;;
      -h|--help)
        usage
        exit 0
        ;;
      *)
        die "unknown option: $1"
        ;;
    esac
  done
}

detect_cpp_sdk_archive() {
  local machine
  machine="$(uname -m)"
  case "${machine}" in
    x86_64|amd64)
      printf 'zeroArm_sdk-linux-x86_64.tar.gz'
      ;;
    aarch64|arm64)
      printf 'zeroArm_sdk-linux-aarch64.tar.gz'
      ;;
    *)
      die "unsupported CPU architecture: ${machine}. Supported: x86_64, aarch64."
      ;;
  esac
}

parse_cpp_release_manifest() {
  python3 - "$1" "$2" <<'PY'
import json
import pathlib
import sys

manifest = json.loads(pathlib.Path(sys.argv[1]).read_text(encoding="utf-8"))
archive_name = sys.argv[2]
matches = [item for item in manifest.get("artifacts", []) if item.get("name") == archive_name]
if len(matches) != 1:
    raise SystemExit(f"manifest must contain exactly one {archive_name} artifact")
artifact = matches[0]
version = str(manifest.get("version", "")).strip()
sha256 = str(artifact.get("sha256", "")).strip().lower()
size = artifact.get("size")
if not version:
    raise SystemExit("manifest version is missing")
if len(sha256) != 64 or any(ch not in "0123456789abcdef" for ch in sha256):
    raise SystemExit("manifest artifact SHA256 is invalid")
if not isinstance(size, int) or size <= 0:
    raise SystemExit("manifest artifact size is invalid")
print(version)
print(sha256)
print(size)
PY
}

validate_workspace() {
  WORKSPACE="$(resolve_dir "${WORKSPACE}")"
  [[ -d "${WORKSPACE}/src/robot_arm_interfaces" ]] || die "missing ${WORKSPACE}/src/robot_arm_interfaces"
  [[ -d "${WORKSPACE}/src/robot_arm_hardware" ]] || die "missing ${WORKSPACE}/src/robot_arm_hardware"
  [[ -d "${WORKSPACE}/src/robot_arm_description" ]] || die "missing ${WORKSPACE}/src/robot_arm_description"
  [[ -d "${WORKSPACE}/src/robot_arm_moveit" ]] || die "missing ${WORKSPACE}/src/robot_arm_moveit"
}

check_tmp_prefix_leak() {
  case ":${CMAKE_PREFIX_PATH:-}:" in
    *":/tmp/zeroarm_cpp_sdk_ros_install:"*)
      die "CMAKE_PREFIX_PATH still contains /tmp/zeroarm_cpp_sdk_ros_install; remove this obsolete test prefix before building."
      ;;
  esac
}

expected_cpp_sdk_dir() {
  printf '%s/lib/cmake/zeroArm_sdk\n' "${CPP_SDK_PREFIX%/}"
}

prepare_cpp_sdk_cache() {
  local cache="${WORKSPACE}/build/robot_arm_hardware/CMakeCache.txt"
  local cached expected
  [[ -f "${cache}" ]] || return 0
  cached="$(sed -n 's/^zeroArm_sdk_DIR:PATH=//p' "${cache}" | head -n 1)"
  expected="$(expected_cpp_sdk_dir)"
  if [[ -n "${cached}" && "${cached}" != "${expected}" ]]; then
    log "C++ SDK prefix changed from ${cached} to ${expected}; clearing robot_arm_hardware CMake cache"
    run rm -rf "${WORKSPACE}/build/robot_arm_hardware"
  fi
}

install_cpp_sdk() {
  if [[ "${SKIP_CPP_SDK}" == "1" ]]; then
    log "Skipping C++ SDK install"
    return 0
  fi

  local archive_name archive_path url installer manifest_path manifest_url
  local version expected_sha256 expected_size actual_size
  local -a release
  archive_name="$(detect_cpp_sdk_archive)"
  url="${CPP_SDK_BASE_URL%/}/${archive_name}"
  manifest_url="${CPP_SDK_BASE_URL%/}/manifest.json"

  log "Installing C++ SDK from ${url}"
  TMP_DIR="$(mktemp -d)"
  archive_path="${TMP_DIR}/${archive_name}"
  manifest_path="${TMP_DIR}/manifest.json"

  run curl -fL --retry 3 --connect-timeout 10 -o "${manifest_path}" "${manifest_url}"
  run curl -fL --retry 3 --connect-timeout 10 -o "${archive_path}" "${url}"

  if [[ "${DRY_RUN}" == "1" ]]; then
    log "Would verify ${archive_name} size and SHA256 against ${manifest_url}"
    run tar -xzf "${archive_path}" -C "${TMP_DIR}"
    run bash "${TMP_DIR}/zeroArm_sdk/install_zeroarm_sdk.sh" manifest-version "${CPP_SDK_PREFIX}"
    return 0
  fi

  mapfile -t release < <(parse_cpp_release_manifest "${manifest_path}" "${archive_name}")
  [[ "${#release[@]}" -eq 3 ]] || die "failed to parse C++ SDK release manifest"
  version="${release[0]}"
  expected_sha256="${release[1]}"
  expected_size="${release[2]}"
  actual_size="$(stat -c '%s' "${archive_path}")"
  [[ "${actual_size}" == "${expected_size}" ]] ||
    die "C++ SDK archive size mismatch: expected ${expected_size}, got ${actual_size}"
  printf '%s  %s\n' "${expected_sha256}" "${archive_name}" |
    (cd -- "${TMP_DIR}" && sha256sum -c -) ||
    die "C++ SDK archive checksum verification failed"
  log "Verified C++ SDK ${version}: ${expected_sha256}"

  tar -xzf "${archive_path}" -C "${TMP_DIR}"
  installer="${TMP_DIR}/zeroArm_sdk/install_zeroarm_sdk.sh"

  [[ -f "${installer}" ]] || die "C++ SDK installer not found inside ${archive_name}"

  if [[ -w "${CPP_SDK_PREFIX}" || ( ! -e "${CPP_SDK_PREFIX}" && -w "$(dirname -- "${CPP_SDK_PREFIX}")" ) ]]; then
    ZEROARM_SKIP_DEPS="${ZEROARM_SKIP_DEPS:-0}" bash "${installer}" "${version}" "${CPP_SDK_PREFIX}"
    return 0
  fi

  command -v sudo >/dev/null 2>&1 || die "prefix ${CPP_SDK_PREFIX} is not writable and sudo is unavailable"
  log "Requesting sudo for C++ SDK install prefix ${CPP_SDK_PREFIX}"
  sudo -v
  sudo env ZEROARM_SKIP_DEPS="${ZEROARM_SKIP_DEPS:-0}" bash "${installer}" "${version}" "${CPP_SDK_PREFIX}"
}

source_ros() {
  local setup="/opt/ros/${ROS_DISTRO}/setup.bash"
  if [[ "${DRY_RUN}" == "1" ]]; then
    log "Would source ${setup}"
    return 0
  fi

  [[ -r "${setup}" ]] || die "ROS setup file not found: ${setup}"
  # shellcheck disable=SC1090
  set +u
  source "${setup}"
  set -u
}

install_ros_dependencies() {
  if [[ "${SKIP_ROSDEP}" == "1" ]]; then
    log "Skipping rosdep install"
    return 0
  fi

  if [[ "${DRY_RUN}" != "1" ]]; then
    command -v rosdep >/dev/null 2>&1 || die "rosdep not found. Install python3-rosdep or rerun with --skip-rosdep."
  fi

  (
    cd -- "${WORKSPACE}"
    run rosdep install --from-paths src --ignore-src -r -y
  )
}

build_workspace() {
  if [[ "${DRY_RUN}" != "1" ]]; then
    command -v colcon >/dev/null 2>&1 || die "colcon not found. Install python3-colcon-common-extensions."
  fi

  (
    cd -- "${WORKSPACE}"
    run env "CMAKE_PREFIX_PATH=${CPP_SDK_PREFIX}:${CMAKE_PREFIX_PATH:-}" \
      colcon build \
      --packages-select "${PACKAGES[@]}" \
      --cmake-args "-DCMAKE_BUILD_TYPE=${BUILD_TYPE}"
  )
}

verify_build_output() {
  if [[ "${DRY_RUN}" == "1" ]]; then
    return 0
  fi

  local cache sdk_dir
  cache="${WORKSPACE}/build/robot_arm_hardware/CMakeCache.txt"
  [[ -f "${cache}" ]] || die "robot_arm_hardware CMake cache was not generated"

  if grep -q '/tmp/zeroarm_cpp_sdk_ros_install' "${cache}"; then
    die "robot_arm_hardware build still references /tmp/zeroarm_cpp_sdk_ros_install"
  fi

  sdk_dir="$(grep -E '^zeroArm_sdk_DIR:PATH=' "${cache}" || true)"
  [[ -n "${sdk_dir}" ]] || die "robot_arm_hardware did not resolve zeroArm_sdk_DIR"
  [[ "${sdk_dir}" == "zeroArm_sdk_DIR:PATH=$(expected_cpp_sdk_dir)" ]] ||
    die "robot_arm_hardware resolved unexpected SDK path: ${sdk_dir}"
  log "Verified ${sdk_dir}"
}

main() {
  parse_args "$@"
  bootstrap_workspace
  if [[ "${DRY_RUN}" == "1" ]] && ! workspace_has_sources; then
    log "Dry run completed before workspace validation"
    return 0
  fi
  validate_workspace
  check_tmp_prefix_leak

  log "Workspace: ${WORKSPACE}"
  log "ROS distro: ${ROS_DISTRO}"
  log "C++ SDK prefix: ${CPP_SDK_PREFIX}"

  install_cpp_sdk
  prepare_cpp_sdk_cache
  source_ros
  install_ros_dependencies
  build_workspace
  verify_build_output

  log "Done. Run: source ${WORKSPACE}/install/setup.bash"
}

main "$@"