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

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)
package_root=$(cd "${script_dir}/.." && pwd -P)
vendor_source="${package_root}/vendor/mtranserver"
build_dir="${package_root}/.build/mtranserver"
runtime_data_build_dir="${package_root}/.build/runtime-data"
runtime_bin_dir="${package_root}/runtime/bin"
runtime_config_dir="${package_root}/runtime/config"
model_root="${package_root}/models"
model_dir="${model_root}/en_zh-Hans"
tools_dir="${script_dir}/tools"
verify_script="${script_dir}/verify_runtime_data.py"
bun_bin=""

BUN_VERSION=1.3.14
os_release_file=${MTRAN_OS_RELEASE_FILE:-/etc/os-release}
machine=${MTRAN_MACHINE:-$(uname -m)}

die()
{
  echo "error: $*" >&2
  exit 1
}

validate_host()
{
  [[ -r "${os_release_file}" ]] || die "cannot read OS release file: ${os_release_file}"
  # shellcheck disable=SC1090
  source "${os_release_file}"
  [[ "${ID:-}" == "ubuntu" && "${VERSION_ID:-}" == "22.04" ]] || \
    die "Ubuntu 22.04 is required."

  case "${machine}" in
    x86_64 | aarch64) ;;
    *) die "supported CPU architecture must be x86_64 or aarch64; found ${machine}." ;;
  esac

  [[ "${ROS_DISTRO:-}" == "humble" ]] || \
    die "ROS2 Humble environment must be sourced before running this script."

  [[ -f "${vendor_source}/package.json" && -f "${vendor_source}/bun.lock" ]] || \
    die "vendored MTranServer source is missing under vendor/mtranserver."
  [[ -f "${verify_script}" ]] || die "runtime verifier is missing: ${verify_script}"
}

validate_ros_packages()
{
  local required_packages=(
    ros-humble-ament-cmake
    ros-humble-ament-cmake-gtest
    ros-humble-rclcpp
    ros-humble-rosidl-default-generators
    ros-humble-rosidl-default-runtime
    ros-humble-launch
    ros-humble-launch-ros
    ros-humble-ament-index-python
  )
  local missing=()
  local package
  for package in "${required_packages[@]}"; do
    if ! dpkg-query -W -f='${Status}' "${package}" 2>/dev/null | grep -Fq 'install ok installed'; then
      missing+=("${package}")
    fi
  done
  if ((${#missing[@]})); then
    printf 'error: missing ROS2 Humble packages (install these separately):\n' >&2
    printf '  %s\n' "${missing[@]}" >&2
    exit 1
  fi
}

install_system_dependencies()
{
  local apt_packages=(
    build-essential
    cmake
    pkg-config
    curl
    ca-certificates
    file
    rsync
    npm
    libcurl4-openssl-dev
    nlohmann-json3-dev
  )
  local missing=()
  local package
  for package in "${apt_packages[@]}"; do
    local available=false
    case "${package}" in
      build-essential)
        command -v c++ >/dev/null 2>&1 && command -v make >/dev/null 2>&1 && available=true
        ;;
      ca-certificates)
        [[ -r /etc/ssl/certs/ca-certificates.crt ]] && available=true
        ;;
      libcurl4-openssl-dev | nlohmann-json3-dev)
        if dpkg-query -W -f='${Status}' "${package}" 2>/dev/null | grep -Fq 'install ok installed'; then
          available=true
        fi
        ;;
      *)
        command -v "${package}" >/dev/null 2>&1 && available=true
        ;;
    esac
    if [[ "${available}" != true ]]; then
      missing+=("${package}")
    fi
  done
  if ((${#missing[@]} == 0)); then
    echo "Reusing installed non-ROS system dependencies."
    return
  fi
  printf 'Installing missing non-ROS system dependencies: %s\n' "${missing[*]}"
  sudo apt-get update
  sudo apt-get install -y "${missing[@]}"
}

prepare_bun()
{
  bun_bin="${tools_dir}/node_modules/.bin/bun"
  export BUN_INSTALL_CACHE_DIR="${tools_dir}/bun-cache"
  export npm_config_cache="${tools_dir}/npm-cache"
  export ELECTRON_SKIP_BINARY_DOWNLOAD=1

  if [[ ! -x "${bun_bin}" || "$("${bun_bin}" --version 2>/dev/null || true)" != "${BUN_VERSION}" ]]; then
    mkdir -p "${tools_dir}"
    npm install --prefix "${tools_dir}" --no-save --no-package-lock "bun@${BUN_VERSION}"
  fi
  [[ -x "${bun_bin}" ]] || die "workspace-local Bun installation failed."
  [[ "$("${bun_bin}" --version)" == "${BUN_VERSION}" ]] || \
    die "workspace-local Bun version must be ${BUN_VERSION}."
  export PATH="$(dirname "${bun_bin}"):${PATH}"
}

sidecar_architecture_matches()
{
  local executable=$1
  local description
  [[ -x "${executable}" ]] || return 1
  description=$(file -b "${executable}" 2>/dev/null) || return 1
  case "${machine}" in
    x86_64) [[ "${description}" == *"x86-64"* || "${description}" == *"x86_64"* ]] ;;
    aarch64) [[ "${description}" == *"aarch64"* || "${description}" == *"ARM64"* ]] ;;
  esac
}

build_sidecar_if_needed()
{
  local bun_bin=$1
  local sidecar="${runtime_bin_dir}/mtranserver"
  local staged_sidecar="${runtime_bin_dir}/.mtranserver.tmp.$$"

  if sidecar_architecture_matches "${sidecar}"; then
    echo "Reusing native MTranServer sidecar: ${sidecar}"
    return
  fi

  mkdir -p "${build_dir}" "${runtime_bin_dir}"
  rsync -a --delete \
    --exclude '.git/' \
    --exclude 'node_modules/' \
    --exclude 'dist/' \
    --exclude 'release/' \
    --exclude '.cache/' \
    "${vendor_source}/" "${build_dir}/"

  (
    cd "${build_dir}"
    "${bun_bin}" install --frozen-lockfile
    "${bun_bin}" run build
  )

  [[ -x "${build_dir}/dist/mtranserver" ]] || die "MTranServer build did not produce dist/mtranserver."
  install -m 0755 "${build_dir}/dist/mtranserver" "${staged_sidecar}"
  if ! sidecar_architecture_matches "${staged_sidecar}"; then
    rm -f "${staged_sidecar}"
    die "built MTranServer architecture does not match ${machine}."
  fi
  mv -f "${staged_sidecar}" "${sidecar}"
  echo "Staged native MTranServer sidecar: ${sidecar}"
}

runtime_data_is_valid()
{
  python3 "${verify_script}" \
    --records "${runtime_config_dir}/records.json" \
    --model-dir "${model_dir}" >/dev/null 2>&1
}

install_verified_runtime_data()
{
  local staged_records="${runtime_data_build_dir}/config/records.json"
  local staged_models="${runtime_data_build_dir}/models/en_zh-Hans"
  local next_models="${model_root}/.en_zh-Hans.next.$$"
  local old_models="${model_root}/.en_zh-Hans.old.$$"
  local next_records="${runtime_config_dir}/.records.json.next.$$"

  python3 "${verify_script}" --records "${staged_records}" --model-dir "${staged_models}"

  mkdir -p "${model_root}" "${runtime_config_dir}"
  rm -rf "${next_models}" "${old_models}"
  mkdir -p "${next_models}"
  cp -a "${staged_models}/." "${next_models}/"

  if [[ -e "${model_dir}" ]]; then
    mv "${model_dir}" "${old_models}"
  fi
  if ! mv "${next_models}" "${model_dir}"; then
    [[ -e "${old_models}" ]] && mv "${old_models}" "${model_dir}"
    die "failed to install verified model directory."
  fi
  rm -rf "${old_models}"

  install -m 0644 "${staged_records}" "${next_records}"
  mv -f "${next_records}" "${runtime_config_dir}/records.json"
}

prepare_runtime_data_if_needed()
{
  local sidecar="${runtime_bin_dir}/mtranserver"
  if runtime_data_is_valid; then
    echo "Reusing verified en -> zh-Hans runtime data."
    return
  fi

  rm -rf "${runtime_data_build_dir}"
  mkdir -p "${runtime_data_build_dir}/config" "${runtime_data_build_dir}/models"
  MT_HOST=127.0.0.1 \
  MT_PORT=0 \
  MT_ENABLE_UI=false \
  MT_OFFLINE=false \
  MT_CHECK_UPDATE=false \
  MT_LOG_LEVEL=info \
  MT_CONFIG_DIR="${runtime_data_build_dir}/config" \
  MT_MODEL_DIR="${runtime_data_build_dir}/models" \
    "${sidecar}" --download en-zh

  install_verified_runtime_data
  runtime_data_is_valid || die "installed runtime data failed final verification."
  echo "Staged verified en -> zh-Hans runtime data."
}

main()
{
  validate_host
  validate_ros_packages
  install_system_dependencies
  prepare_bun
  build_sidecar_if_needed "${bun_bin}"
  prepare_runtime_data_if_needed

  local workspace_root
  workspace_root=$(cd "${package_root}/../.." && pwd -P)
  printf 'Preparation complete.\n'
  printf 'Next: cd %q && colcon build --packages-select mtran\n' "${workspace_root}"
  printf 'Then: source install/setup.bash && ros2 launch mtran translator.launch.py\n'
}

main "$@"
