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

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# shellcheck source=./lib.sh
source "$SCRIPT_DIR/lib.sh"

TIMEZONE="America/Chicago"
SSH_PORT="22"
AUTO_YES=false
ENABLE_UFW=true

usage() {
  cat <<USAGE
Usage: $(basename "$0") [options]

Prepare a brand-new Raspberry Pi OS host for monitoring stack workloads.

Options:
  --timezone=ZONE      Set system timezone (default: America/Chicago)
  --ssh-port=PORT      SSH port allowed by firewall (default: 22)
  --skip-firewall      Skip UFW configuration
  --yes, -y            Non-interactive; skip confirmation prompts
  --help, -h           Show help

Example:
  $(basename "$0") --timezone=America/Chicago --yes
USAGE
}

for arg in "$@"; do
  case "$arg" in
    --timezone=*) TIMEZONE="${arg#*=}" ;;
    --ssh-port=*) SSH_PORT="${arg#*=}" ;;
    --skip-firewall) ENABLE_UFW=false ;;
    --yes|-y) AUTO_YES=true ;;
    --help|-h) usage; exit 0 ;;
    *) log_error "Unknown argument: $arg"; usage; exit 1 ;;
  esac
done

# Validate --ssh-port (must be 1-65535) before we risk enabling UFW with a bad rule
if ! [[ "$SSH_PORT" =~ ^[0-9]+$ ]] || [[ "$SSH_PORT" -lt 1 ]] || [[ "$SSH_PORT" -gt 65535 ]]; then
  log_error "--ssh-port must be a number between 1 and 65535 (got: '$SSH_PORT')"
  exit 1
fi

# Validate --timezone against timedatectl's known list
if ! timedatectl list-timezones 2>/dev/null | grep -qx "$TIMEZONE"; then
  log_error "Unknown timezone: '$TIMEZONE'"
  log_error "Run 'timedatectl list-timezones' for valid options"
  exit 1
fi

require_cmd sudo apt systemctl timedatectl curl

if ! confirm_action "This will install/update OS packages and Docker on this Pi. Continue?" "$AUTO_YES"; then
  log_info "Cancelled"
  exit 0
fi

arch="$(uname -m)"
if [[ "$arch" != "aarch64" && "$arch" != "arm64" ]]; then
  log_warn "Detected architecture '$arch' (expected ARM64 for Raspberry Pi 4)."
fi

log_info "Updating OS packages..."
sudo apt update
sudo apt full-upgrade -y

log_info "Installing base packages..."
sudo apt install -y \
  ca-certificates \
  curl \
  fail2ban \
  git \
  htop \
  jq \
  ufw \
  unattended-upgrades

log_info "Setting timezone to $TIMEZONE..."
sudo timedatectl set-timezone "$TIMEZONE"

log_info "Enabling security services..."
sudo systemctl enable --now fail2ban
sudo systemctl enable --now unattended-upgrades || true

if ! command -v docker >/dev/null 2>&1; then
  log_info "Installing Docker Engine..."
  curl -fsSL https://get.docker.com | sh
else
  log_success "Docker already installed"
fi

log_info "Enabling Docker service..."
sudo systemctl enable --now docker

log_info "Configuring Docker daemon defaults..."
sudo mkdir -p /etc/docker
if [[ -f /etc/docker/daemon.json ]]; then
  sudo cp /etc/docker/daemon.json /etc/docker/daemon.json.bak
  log_info "Backed up existing daemon.json to daemon.json.bak"
fi
sudo tee /etc/docker/daemon.json >/dev/null <<'JSON'
{
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "10m",
    "max-file": "3"
  },
  "live-restore": true
}
JSON
sudo systemctl restart docker

if id -nG "$USER" | grep -qw docker; then
  log_success "User '$USER' is already in docker group"
else
  log_info "Adding '$USER' to docker group..."
  sudo usermod -aG docker "$USER"
  log_warn "Log out and back in once for docker group membership to apply."
fi

if [[ "$ENABLE_UFW" == "true" ]]; then
  log_info "Configuring UFW firewall rules..."
  sudo ufw allow "${SSH_PORT}/tcp"
  sudo ufw allow 9443/tcp   # Portainer HTTPS
  sudo ufw allow 8000/tcp   # Portainer Edge
  sudo ufw allow 3000/tcp   # Grafana
  sudo ufw allow 3001/tcp   # Uptime Kuma
  sudo ufw allow 9090/tcp   # Prometheus
  sudo ufw --force enable
fi

log_success "Bootstrap complete"
log_info "Recommended next steps:"
log_info "1) Re-login to apply docker group membership"
log_info "2) (Optional) Run setup/pi-monitoring/mount_ssd.sh if you have an SSD"
log_info "3) Copy stack.env.example to stack.env and run deploy_stack.sh"