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

bundle_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
zero_dir="$bundle_dir/zero"
image_tar="$zero_dir/product_vamp_zero_linux_arm64.docker.tar"
config_dir=/etc/product_vamp
state_dir=/var/lib/product_vamp
run_dir=/run/product_vamp
share_dir=/usr/share/product_vamp
image_tag=
connection=product_vamp_control
replacement_started=false
image_loaded=false
previous_image=
stage_dir=
legacy_targets=(
  /etc/product-vamp /var/lib/product-vamp /run/product-vamp /usr/share/product-vamp
  /usr/local/bin/product-vamp-start-ap /usr/local/bin/product-vamp-start-zero
  /etc/systemd/system/product-vamp-ap.service /etc/systemd/system/product-vamp-zero.service
  /etc/systemd/system/product-vamp-ap.service.d /etc/systemd/system/product-vamp-zero.service.d
  /run/systemd/system/product-vamp-ap.service.d /run/systemd/system/product-vamp-zero.service.d
)

cleanup() {
  status=$?
  trap - EXIT
  if [ "$status" -ne 0 ] && [ "$replacement_started" = true ]; then
    systemctl stop product_vamp_zero.service product_vamp_ap.service >/dev/null 2>&1 || true
    docker rm -f product_vamp_zero >/dev/null 2>&1 || true
    nmcli connection down "$connection" >/dev/null 2>&1 || true
    echo "Installation failed; Product Vamp services are stopped. Rerun the installer." >&2
  elif [ "$status" -ne 0 ] && [ "$image_loaded" = true ] && [ -n "$previous_image" ]; then
    docker tag "$previous_image" "$image_tag" || echo "Failed to restore the previous image tag." >&2
  fi
  if [ -n "$stage_dir" ]; then
    rm -rf -- "$stage_dir"
  fi
  exit "$status"
}
trap cleanup EXIT

validate_target() {
  target=$1
  if [ "$(realpath -m -- "$target")" != "$target" ] || [ -L "$target" ]; then
    echo "Refusing unsafe installation target: $target" >&2
    exit 1
  fi
  if [ -e "$target" ] && [ -d "$target" ]; then
    if [ -n "$(find "$target" -type l -print -quit)" ]; then
      echo "Refusing installation target containing symlinks: $target" >&2
      exit 1
    fi
    mounts=$(find "$target" -xdev -type d -exec mountpoint -q {} ';' -print)
    if [ -n "$mounts" ]; then
      echo "Refusing mounted installation target: $target" >&2
      exit 1
    fi
  fi
}
required_packages=(
  docker.io
  docker-cli
  network-manager
  sqlite3
  ca-certificates
  python3
)

if [ "$(id -u)" -ne 0 ]; then
  echo "install_zero.sh must be run as root" >&2
  exit 1
fi

if [ ! -f "$image_tar" ]; then
  echo "missing Docker archive: $image_tar" >&2
  exit 1
fi

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

if [ "${#missing_packages[@]}" -eq 0 ]; then
  echo "Required OS packages already installed; skipping apt."
else
  echo "Installing missing OS packages: ${missing_packages[*]}"
  if ! apt-get update -o APT::Update::Error-Mode=any; then
    echo "Failed to refresh apt package indexes." >&2
    echo "Connect the Zero to upstream internet, then rerun sudo zero/install_zero.sh." >&2
    exit 1
  fi
  if ! apt-get install -y --no-install-recommends "${missing_packages[@]}"; then
    echo "Failed to install required OS packages: ${missing_packages[*]}" >&2
    echo "Connect the Zero to upstream internet, then rerun sudo zero/install_zero.sh." >&2
    exit 1
  fi
fi
systemctl enable --now docker NetworkManager

image_details=$(python3 - "$bundle_dir/manifest.json" "$image_tar" <<'PY'
import hashlib
import json
import re
import sys
import tarfile
with open(sys.argv[1], encoding='utf-8') as source:
  tag = json.load(source)['zero']['image_tag']
if not isinstance(tag, str) or not re.fullmatch(r'[A-Za-z0-9][A-Za-z0-9._:/-]*', tag):
  raise ValueError('invalid bundle image tag')
with tarfile.open(sys.argv[2]) as archive:
  with archive.extractfile('manifest.json') as source:
    images = [image for image in json.load(source) if tag in (image.get('RepoTags') or [])]
  if len(images) != 1:
    raise ValueError('Docker archive does not contain exactly one matching bundle image')
  with archive.extractfile(images[0]['Config']) as source:
    digest = hashlib.file_digest(source, 'sha256').hexdigest()
print(tag)
print('sha256:' + digest)
PY
)
mapfile -t image_fields <<< "$image_details"
image_tag=${image_fields[0]}
expected_image=${image_fields[1]}

for target in "${legacy_targets[@]}" "$config_dir" "$state_dir" "$run_dir" "$share_dir" \
  /usr/local/bin/product_vamp_start_ap /usr/local/bin/product_vamp_start_zero \
  /etc/systemd/system/product_vamp_ap.service.d /etc/systemd/system/product_vamp_zero.service.d \
  /run/systemd/system/product_vamp_ap.service.d /run/systemd/system/product_vamp_zero.service.d \
  /etc/systemd/system/product_vamp_ap.service /etc/systemd/system/product_vamp_zero.service; do
  validate_target "$target"
done

stage_dir=$(mktemp -d /tmp/product_vamp_install_XXXXXX)
install -m 0644 "$zero_dir/config/vamp.example.json" "$stage_dir/vamp.json"
install -m 0755 "$zero_dir/start_ap.sh" "$stage_dir/start_ap.sh"
install -m 0755 "$zero_dir/start_zero.sh" "$stage_dir/start_zero.sh"
install -m 0644 "$zero_dir/systemd/product_vamp_ap.service" "$stage_dir/product_vamp_ap.service"
install -m 0644 "$zero_dir/systemd/product_vamp_zero.service" "$stage_dir/product_vamp_zero.service"
bash -n "$stage_dir/start_ap.sh"
bash -n "$stage_dir/start_zero.sh"
previous_image=$(docker image inspect --format '{{.Id}}' "$image_tag" 2>/dev/null || true)
previous_images=$(docker image ls --quiet --no-trunc --filter label=org.khodo.product=product_vamp_zero)
legacy_images=$(docker image ls --quiet --no-trunc --filter label=org.khodo.product=product-vamp-zero)
image_loaded=true
docker load -i "$image_tar"
new_image=$(docker image inspect --format '{{.Id}}' "$image_tag")
if [ "$new_image" != "$expected_image" ]; then
  echo "Loaded Docker image does not match the incoming archive." >&2
  exit 1
fi
# Probe the actual image and staged configuration without touching host state or hardware.
docker run --rm --network none --mount "type=bind,src=$stage_dir,dst=/install,readonly" \
  "$new_image" /usr/bin/vamp_zero --config /install/vamp.json --once

echo "Resetting all Product Vamp data and settings, including control Wi-Fi credentials."
echo "A connection through the control Wi-Fi may disconnect."
replacement_started=true
for service in product_vamp_zero.service product_vamp_ap.service product-vamp-zero.service product-vamp-ap.service; do
  if systemctl cat "$service" >/dev/null 2>&1; then
    systemctl stop "$service"
    systemctl disable "$service"
  fi
done
for container in product_vamp_zero product-vamp-zero; do
  if docker container inspect "$container" >/dev/null 2>&1; then
    docker rm -f "$container"
  fi
done
for target in "${legacy_targets[@]}" "$config_dir" "$state_dir" "$run_dir" "$share_dir" \
  /etc/systemd/system/product_vamp_ap.service.d /etc/systemd/system/product_vamp_zero.service.d \
  /run/systemd/system/product_vamp_ap.service.d /run/systemd/system/product_vamp_zero.service.d; do
  validate_target "$target"
  rm -rf --one-file-system -- "$target"
done
install -d -m 0755 "$config_dir" "$run_dir"
install -m 0644 "$stage_dir/vamp.json" "$config_dir/vamp.json"
printf '%s\n' "$new_image" > "$config_dir/image_id"
install -m 0755 "$stage_dir/start_ap.sh" /usr/local/bin/product_vamp_start_ap
install -m 0755 "$stage_dir/start_zero.sh" /usr/local/bin/product_vamp_start_zero
install -m 0644 "$stage_dir/product_vamp_ap.service" /etc/systemd/system/product_vamp_ap.service
install -m 0644 "$stage_dir/product_vamp_zero.service" /etc/systemd/system/product_vamp_zero.service

for profile in "$connection" product-vamp-control; do
  if nmcli -t -f NAME connection show | grep -Fxq "$profile"; then
    nmcli connection delete "$profile"
  fi
done
systemctl daemon-reload
systemctl disable --now getty@tty1.service >/dev/null 2>&1 || true
systemctl enable --now getty@tty2.service >/dev/null 2>&1 || true
systemctl enable --now product_vamp_ap.service
systemctl enable product_vamp_zero.service
systemctl restart product_vamp_zero.service
healthy=false
for attempt in $(seq 1 60); do
  if docker exec product_vamp_zero /usr/bin/vamp_zero --config /etc/product_vamp/vamp.json --health-check; then
    healthy=true
    break
  fi
  sleep 1
done
if [ "$healthy" != true ]; then
  echo "Product Vamp did not become healthy after installation." >&2
  exit 1
fi
for image in $(printf '%s\n' "$previous_image" "$previous_images" "$legacy_images" | sort -u); do
  if [ "$image" != "$new_image" ]; then
    if [ -z "$(docker ps -aq --filter "ancestor=$image")" ]; then
      docker image rm "$image" || echo "Warning: could not remove old Product Vamp image $image" >&2
    else
      echo "Keeping old Product Vamp image $image: another container still uses it." >&2
    fi
  fi
done
echo "Product Vamp Zero installed. Control AP defaults to SSID 'Vamp Control' at http://192.168.42.1:8080."