#!/usr/bin/env bash
# Stage all network-fetched Server Starter inputs for a disconnected install.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PACKAGE_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
OUTPUT_DIR="${OFFLINE_OUTPUT_DIR:-$PACKAGE_DIR/autoinstall/offline-output}"
PAYLOAD_DIR="$OUTPUT_DIR/payload"
MANIFEST_DIR="$OUTPUT_DIR/manifest"

# shellcheck source=payload.env
source "$SCRIPT_DIR/payload.env"

if [ "${EUID:-$(id -u)}" -ne 0 ]; then
  echo "ERROR: run as root on a connected Ubuntu 22.04 staging host." >&2
  exit 1
fi
. /etc/os-release
if [ "${ID:-}" != ubuntu ] || [ "${VERSION_ID:-}" != 22.04 ]; then
  echo "ERROR: Ubuntu 22.04 staging host required; found ${PRETTY_NAME:-unknown}." >&2
  exit 1
fi

for command in apt-get curl docker git python3 rsync sha256sum tar; do
  command -v "$command" >/dev/null || {
    echo "ERROR: missing staging command: $command" >&2
    exit 1
  }
done

available_kb="$(df -Pk "$(dirname "$OUTPUT_DIR")" | awk 'NR==2 {print $4}')"
if [ "$available_kb" -lt $((80 * 1024 * 1024)) ]; then
  echo "ERROR: at least 80 GiB free is required for the Workstation offline payload." >&2
  exit 1
fi

rm -rf "$OUTPUT_DIR"
mkdir -p "$PAYLOAD_DIR"/{apt,containers,models,runtime,source} \
  "$PAYLOAD_DIR/python"/{backend,chromadb,platform} "$MANIFEST_DIR"

echo "-> Capturing pinned appliance source"
if git -C "$PACKAGE_DIR" rev-parse HEAD >/dev/null 2>&1; then
  git -C "$PACKAGE_DIR" rev-parse HEAD > "$MANIFEST_DIR/source-commit.txt"
elif [ -n "${OFFLINE_SOURCE_COMMIT:-}" ]; then
  printf '%s\n' "$OFFLINE_SOURCE_COMMIT" > "$MANIFEST_DIR/source-commit.txt"
else
  echo "ERROR: source commit is unavailable; set OFFLINE_SOURCE_COMMIT for exported source trees." >&2
  exit 1
fi
rsync -a --delete \
  --exclude '.git/' --exclude '*.iso' --exclude 'offline-output/' \
  --exclude 'how-it-flows.gif' \
  "$PACKAGE_DIR/" "$PAYLOAD_DIR/source/aipackage/"

echo "-> Downloading Ubuntu package payload"
export DEBIAN_FRONTEND=noninteractive
mapfile -t apt_packages < <(sed -E '/^[[:space:]]*(#|$)/d' "$SCRIPT_DIR/apt-packages.txt")
apt-get update
apt-get install -y --download-only --reinstall \
  -o "Dir::Cache::archives=$PAYLOAD_DIR/apt" "${apt_packages[@]}"

echo "-> Downloading Python wheelhouse"
python3 -m pip download --dest "$PAYLOAD_DIR/python/backend" \
  -r "$PACKAGE_DIR/ansible/roles/cezen-backend/files/requirements.txt"
python3 -m pip download --dest "$PAYLOAD_DIR/python/chromadb" \
  'chromadb==0.5.23'
grep -vE '^[[:space:]]*(#|$|chromadb==)' "$SCRIPT_DIR/python-packages.txt" \
  > "$OUTPUT_DIR/platform-python-packages.txt"
python3 -m pip download --dest "$PAYLOAD_DIR/python/platform" \
  -r "$OUTPUT_DIR/platform-python-packages.txt"
rm -f "$OUTPUT_DIR/platform-python-packages.txt"

echo "-> Capturing Miniconda installer"
curl --fail --location --retry 3 "$MINICONDA_URL" \
  --output "$PAYLOAD_DIR/runtime/miniconda.sh"
chmod 0755 "$PAYLOAD_DIR/runtime/miniconda.sh"

echo "-> Pulling and exporting container images"
: > "$MANIFEST_DIR/container-images.txt"
while IFS= read -r image; do
  case "$image" in ''|'#'*) continue ;; esac
  docker pull "$image"
  safe_name="$(printf '%s' "$image" | tr '/:@' '___')"
  docker save "$image" | gzip -1 > "$PAYLOAD_DIR/containers/$safe_name.tar.gz"
  digest="$(docker image inspect "$image" --format '{{join .RepoDigests ","}}')"
  printf '%s\t%s\n' "$image" "$digest" >> "$MANIFEST_DIR/container-images.txt"
done < "$SCRIPT_DIR/container-images.txt"

echo "-> Capturing Ollama runtime and Starter models"
if ! command -v ollama >/dev/null; then
  curl --fail --location --retry 3 https://ollama.com/install.sh | sh
fi
tar -C / -czf "$PAYLOAD_DIR/runtime/ollama-runtime.tar.gz" \
  usr/local/bin/ollama usr/local/lib/ollama
isolated_ollama="$OUTPUT_DIR/ollama-staging"
mkdir -p "$isolated_ollama"
OLLAMA_MODELS="$isolated_ollama/models" OLLAMA_HOST=127.0.0.1:11435 \
  ollama serve >"$OUTPUT_DIR/ollama-staging.log" 2>&1 &
ollama_pid=$!
cleanup_ollama() { kill "$ollama_pid" 2>/dev/null || true; }
trap cleanup_ollama EXIT
for _ in $(seq 1 30); do
  OLLAMA_HOST=127.0.0.1:11435 ollama list >/dev/null 2>&1 && break
  sleep 1
done
for model in $OLLAMA_MODELS; do
  OLLAMA_MODELS="$isolated_ollama/models" OLLAMA_HOST=127.0.0.1:11435 \
    ollama pull "$model"
done
cleanup_ollama
trap - EXIT
tar -C "$isolated_ollama" -czf "$PAYLOAD_DIR/models/ollama-models.tar.gz" models
rm -rf "$isolated_ollama" "$OUTPUT_DIR/ollama-staging.log"

echo "-> Capturing vLLM Starter model snapshot"
python3 -m pip install --quiet huggingface-hub
python3 - "$VLLM_MODEL" "$PAYLOAD_DIR/models/huggingface" <<'PY'
import sys
from huggingface_hub import snapshot_download
snapshot_download(repo_id=sys.argv[1], local_dir=sys.argv[2])
PY

echo "-> Writing provenance and integrity manifest"
{
  printf 'schema=%s\n' "$OFFLINE_SCHEMA"
  printf 'tier=%s\n' "$OFFLINE_TIER"
  printf 'ubuntu_release=%s\n' "$UBUNTU_RELEASE"
  printf 'built_at=%s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  printf 'builder=%s\n' "$(hostname -f 2>/dev/null || hostname)"
} > "$MANIFEST_DIR/provenance.env"
(
  cd "$OUTPUT_DIR"
  find payload manifest -type f ! -name SHA256SUMS -print0 \
    | sort -z | xargs -0 sha256sum > manifest/SHA256SUMS
  sha256sum -c manifest/SHA256SUMS >/dev/null
)

echo "Offline payload staged at $OUTPUT_DIR"
du -sh "$OUTPUT_DIR"