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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
CRATE_MANIFEST="$ROOT_DIR/traffic_core/Cargo.toml"
TARGET_BIN="$ROOT_DIR/traffic_core/target/release/awg2-harness"
RUST_TOOLCHAIN_FILE="$ROOT_DIR/rust-toolchain.toml"
RUST_TOOLCHAIN_VERSION="1.95.0"

PREPARE_BUILD_ENV=false
PREPARE_RUNTIME_ENV=false
BOOTSTRAP_RUST=false
INSTALL=false
ENABLE=false
YES=false
SKIP_FMT=false

RUN_AS_ROOT=()

BUILD_PACKAGES=(
  ca-certificates
  curl
  build-essential
  pkg-config
  file
  shellcheck
)

RUNTIME_PACKAGES=(
  iproute2
  wireguard-tools
  nftables
  kmod
  util-linux
  procps
  systemd
)

usage() {
  cat <<'EOF'
Usage:
  ./tools/build.sh
  ./tools/build.sh --prepare-build-env --yes
  ./tools/build.sh --prepare-runtime-env --yes
  ./tools/build.sh --bootstrap-rust --yes
  ./tools/build.sh --install
  ./tools/build.sh --enable

Flags:
  --prepare-build-env   Install missing build dependencies via apt
  --prepare-runtime-env Install missing runtime dependencies via apt
  --bootstrap-rust      Install rustup/toolchain pinned by this repo
  --install             Install binary and systemd unit (no enable/start)
  --enable              Enable and start awg2-harness.service
  --yes                 Confirm system mutations (apt/rustup)
  --skip-fmt            Skip cargo fmt --check
  --help                Show this help
EOF
}

for arg in "$@"; do
  case "$arg" in
    --prepare-build-env) PREPARE_BUILD_ENV=true ;;
    --prepare-runtime-env) PREPARE_RUNTIME_ENV=true ;;
    --bootstrap-rust) BOOTSTRAP_RUST=true ;;
    --install) INSTALL=true ;;
    --enable) ENABLE=true ;;
    --yes) YES=true ;;
    --skip-fmt) SKIP_FMT=true ;;
    --help|-h) usage; exit 0 ;;
    *) echo "Unknown arg: $arg"; usage; exit 1 ;;
  esac
done

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || {
    echo "Missing command: $1"
    exit 1
  }
}

setup_root_runner() {
  if [[ "$(id -u)" -eq 0 ]]; then
    RUN_AS_ROOT=()
  else
    require_cmd sudo
    RUN_AS_ROOT=(sudo)
  fi
}

ensure_supported_host() {
  if [[ -f /etc/os-release ]]; then
    # shellcheck disable=SC1091
    source /etc/os-release
    case "${ID:-}" in
      debian|ubuntu) ;;
      *) echo "Unsupported OS: ${ID:-unknown}. Only Debian/Ubuntu are supported."; exit 1 ;;
    esac
  fi

  local arch
  arch="$(uname -m)"
  if [[ "$arch" != "x86_64" && "$arch" != "amd64" ]]; then
    echo "Unsupported arch: $arch. Only amd64/x86_64 is supported."
    exit 1
  fi
}

apt_install_missing() {
  local packages=("$@")
  local missing=()

  require_cmd dpkg-query
  require_cmd apt-get

  for pkg in "${packages[@]}"; do
    if ! dpkg-query -W -f='${Status}' "$pkg" 2>/dev/null | grep -q "install ok installed"; then
      missing+=("$pkg")
    fi
  done

  if [[ "${#missing[@]}" -eq 0 ]]; then
    echo "All required packages already installed"
    return 0
  fi

  echo "Missing packages: ${missing[*]}"
  if [[ "$YES" != true ]]; then
    echo "Re-run with --yes to install them"
    exit 1
  fi

  setup_root_runner
  "${RUN_AS_ROOT[@]}" apt-get update
  "${RUN_AS_ROOT[@]}" apt-get install -y --no-install-recommends "${missing[@]}"
}

bootstrap_rust() {
  require_cmd curl

  if [[ "$YES" != true ]]; then
    echo "This will install rustup and Rust toolchain $RUST_TOOLCHAIN_VERSION"
    echo "Re-run with --yes to proceed"
    exit 1
  fi

  if ! command -v rustup >/dev/null 2>&1; then
    local rustup_script
    rustup_script="$(mktemp)"
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs -o "$rustup_script"
    sh "$rustup_script" -y --profile minimal --default-toolchain "$RUST_TOOLCHAIN_VERSION"
    rm -f "$rustup_script"
  fi

  export PATH="$HOME/.cargo/bin:$PATH"
  require_cmd rustup
  rustup toolchain install "$RUST_TOOLCHAIN_VERSION" --profile minimal --component rustfmt --component clippy
  rustup component add rustfmt clippy --toolchain "$RUST_TOOLCHAIN_VERSION"
  rustc --version
}

prepare_build_gate() {
  require_cmd cargo
  require_cmd rustc
  require_cmd gcc
  require_cmd pkg-config
  require_cmd file
  require_cmd ldd

  if [[ ! -f "$RUST_TOOLCHAIN_FILE" ]]; then
    echo "Missing rust toolchain pin file: $RUST_TOOLCHAIN_FILE"
    exit 1
  fi

  if [[ ! -f "$ROOT_DIR/traffic_core/Cargo.lock" ]]; then
    echo "Missing lockfile: $ROOT_DIR/traffic_core/Cargo.lock"
    exit 1
  fi

  if [[ "$SKIP_FMT" != true ]]; then
    cargo fmt --manifest-path "$CRATE_MANIFEST" --all --check
  fi

  cargo clippy --locked --manifest-path "$CRATE_MANIFEST" --all-targets -- -D warnings
  cargo test --locked --manifest-path "$CRATE_MANIFEST" --release
  cargo build --locked --manifest-path "$CRATE_MANIFEST" --release

  file "$TARGET_BIN"
  ldd "$TARGET_BIN" || true
  "$TARGET_BIN" --help >/dev/null
  sha256sum "$TARGET_BIN"
  ls -lh "$TARGET_BIN"
}

install_artifacts() {
  require_cmd install
  require_cmd systemctl

  if [[ ! -f "$TARGET_BIN" ]]; then
    echo "Binary not found: $TARGET_BIN"
    echo "Run build first: ./tools/build.sh"
    exit 1
  fi

  if [[ ! -f /etc/awg2-harness/config.toml ]]; then
    echo "Missing /etc/awg2-harness/config.toml"
    echo "Refusing to install example config as production config"
    echo "Create config manually from $ROOT_DIR/config.example.toml and re-run with --install"
    exit 1
  fi

  setup_root_runner
  "${RUN_AS_ROOT[@]}" install -m 0755 "$TARGET_BIN" /usr/local/bin/awg2-harness
  "${RUN_AS_ROOT[@]}" install -m 0644 "$ROOT_DIR/awg2-harness.service" /etc/systemd/system/awg2-harness.service

  if command -v systemd-analyze >/dev/null 2>&1; then
    "${RUN_AS_ROOT[@]}" systemd-analyze verify /etc/systemd/system/awg2-harness.service
  fi
  "${RUN_AS_ROOT[@]}" systemctl daemon-reload
}

enable_service() {
  require_cmd systemctl
  setup_root_runner
  "${RUN_AS_ROOT[@]}" systemctl enable --now awg2-harness.service
}

ensure_supported_host

if [[ "$PREPARE_BUILD_ENV" == true ]]; then
  apt_install_missing "${BUILD_PACKAGES[@]}"
fi

if [[ "$PREPARE_RUNTIME_ENV" == true ]]; then
  apt_install_missing "${RUNTIME_PACKAGES[@]}"
fi

if [[ "$BOOTSTRAP_RUST" == true ]]; then
  bootstrap_rust
fi

prepare_build_gate

if [[ "$INSTALL" == true ]]; then
  install_artifacts
fi

if [[ "$ENABLE" == true ]]; then
  enable_service
fi

echo "Build completed: $TARGET_BIN"