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

# -------------------------------------------------------------------
# Mozart bootstrap — run once as root on a fresh Ubuntu 24.04 Droplet.
#
# Usage (with domain — HTTPS):
#   bash deploy/bootstrap.sh \
#     --repo https://github.com/seanli/mozart.git \
#     --key sk-or-xxxx \
#     --domain mozart.example.com
#
# Usage (IP only — HTTP, no TLS):
#   bash deploy/bootstrap.sh \
#     --repo https://github.com/seanli/mozart.git \
#     --key sk-or-xxxx
#
# Private repo (use a GitHub PAT in the URL):
#   bash deploy/bootstrap.sh \
#     --repo https://TOKEN@github.com/seanli/mozart.git \
#     --key sk-or-xxxx
# -------------------------------------------------------------------

REPO_URL=""
API_KEY=""
DOMAIN=""
SSH_PUBKEY=""
MOZART_USER="mozart"
MOZART_HOME="/home/${MOZART_USER}"
REPO_DIR="${MOZART_HOME}/mozart"

usage() {
  cat <<EOF
Usage: bootstrap.sh [OPTIONS]

Required:
  --repo URL        Git repo URL (https or ssh)
  --key KEY         OpenRouter API key

Optional:
  --domain FQDN     Domain name for Caddy TLS (e.g. mozart.example.com)
                    If omitted, serves over HTTP on the Droplet's public IP.
  --ssh-pubkey KEY  SSH public key to add to mozart user's authorized_keys
                    (if omitted, copies root's authorized_keys)
EOF
  exit 1
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --repo)     REPO_URL="$2"; shift 2 ;;
    --key)      API_KEY="$2"; shift 2 ;;
    --domain)   DOMAIN="$2"; shift 2 ;;
    --ssh-pubkey) SSH_PUBKEY="$2"; shift 2 ;;
    *) echo "Unknown option: $1"; usage ;;
  esac
done

[[ -z "$REPO_URL" ]] && { echo "Error: --repo is required"; usage; }
[[ -z "$API_KEY" ]]  && { echo "Error: --key is required"; usage; }

# If no domain provided, auto-detect public IP and serve plain HTTP
if [[ -z "$DOMAIN" ]]; then
  PUBLIC_IP=$(curl -sf https://api.ipify.org || hostname -I | awk '{print $1}')
  CADDY_ADDRESS="http://${PUBLIC_IP}"
  echo "No --domain provided. Using IP-only mode: ${CADDY_ADDRESS}"
  echo "  (basic auth over HTTP — add a domain later for TLS)"
else
  CADDY_ADDRESS="$DOMAIN"
fi

if [[ "$(id -u)" -ne 0 ]]; then
  echo "Error: must run as root"
  exit 1
fi

echo "=== [1/8] Creating user '${MOZART_USER}' ==="
if ! id "$MOZART_USER" &>/dev/null; then
  adduser --disabled-password --gecos "" "$MOZART_USER"
  usermod -aG sudo "$MOZART_USER"
  echo "${MOZART_USER} ALL=(ALL) NOPASSWD:ALL" > "/etc/sudoers.d/${MOZART_USER}"
fi

mkdir -p "${MOZART_HOME}/.ssh"
chmod 700 "${MOZART_HOME}/.ssh"
if [[ -n "$SSH_PUBKEY" ]]; then
  echo "$SSH_PUBKEY" >> "${MOZART_HOME}/.ssh/authorized_keys"
elif [[ -f /root/.ssh/authorized_keys ]]; then
  cp /root/.ssh/authorized_keys "${MOZART_HOME}/.ssh/authorized_keys"
fi
chmod 600 "${MOZART_HOME}/.ssh/authorized_keys"
chown -R "${MOZART_USER}:${MOZART_USER}" "${MOZART_HOME}/.ssh"

echo "=== [2/8] Installing system packages ==="
apt-get update
apt-get install -y --no-install-recommends \
  podman slirp4netns uidmap fuse-overlayfs \
  unzip curl git \
  debian-keyring debian-archive-keyring apt-transport-https

# Caddy from official repo (gets latest stable with systemd integration)
curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' \
  | gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg
curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' \
  | tee /etc/apt/sources.list.d/caddy-stable.list
apt-get update
apt-get install -y caddy

# Podman on Ubuntu doesn't search Docker Hub by default
if [[ ! -f /etc/containers/registries.conf.d/docker.conf ]]; then
  mkdir -p /etc/containers/registries.conf.d
  echo 'unqualified-search-registries = ["docker.io"]' \
    > /etc/containers/registries.conf.d/docker.conf
fi

# Rootless Podman: allow the mozart user's containers to survive logout
loginctl enable-linger "$MOZART_USER"

echo "=== [3/8] Creating 2 GB swap file ==="
# 2GB swap is needed for the Vite UI build (~6700 modules), which peaks above 2GB RAM.
# At runtime, agents are I/O-bound and swap is rarely touched.
if [[ ! -f /swapfile ]]; then
  fallocate -l 2G /swapfile
  chmod 600 /swapfile
  mkswap /swapfile
  swapon /swapfile
  echo '/swapfile none swap sw 0 0' >> /etc/fstab
  echo "Swap enabled (2 GB)."
else
  echo "Swap already exists, skipping."
fi

echo "=== [4/8] Configuring firewall ==="
ufw --force reset
ufw default deny incoming
ufw default allow outgoing
ufw allow ssh
ufw allow 80/tcp
ufw allow 443/tcp
ufw --force enable

# Harden SSH
sed -i 's/^#\?PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
sed -i 's/^#\?PermitRootLogin.*/PermitRootLogin no/' /etc/ssh/sshd_config
systemctl reload ssh

echo "=== [5/8] Installing Bun ==="
su - "$MOZART_USER" -c 'curl -fsSL https://bun.sh/install | bash'

echo "=== [6/8] Cloning repo and building ==="
su - "$MOZART_USER" <<BUILDEOF
  set -euo pipefail
  export PATH="\$HOME/.bun/bin:\$PATH"

  if [[ -d "${REPO_DIR}" ]]; then
    cd "${REPO_DIR}" && git pull --ff-only
  else
    git clone "${REPO_URL}" "${REPO_DIR}"
    cd "${REPO_DIR}"
  fi

  bun install
  cd ui && bun install && cd ..
  bun run build
BUILDEOF

echo "=== [7/8] Running mozart setup (Podman images) ==="
su - "$MOZART_USER" <<SETUPEOF
  set -euo pipefail
  export PATH="\$HOME/.bun/bin:\$PATH"
  cd "${REPO_DIR}"
  ./mozart setup
SETUPEOF

# Write env file for systemd (printf avoids interpreting special chars in the key)
mkdir -p "${MOZART_HOME}/.mozart"
printf 'OPENROUTER_API_KEY=%s\n' "$API_KEY" > "${MOZART_HOME}/.mozart/env"
chmod 600 "${MOZART_HOME}/.mozart/env"
chown -R "${MOZART_USER}:${MOZART_USER}" "${MOZART_HOME}/.mozart"

echo "=== [8/8] Installing systemd service and Caddy config ==="

# Resolve the actual UID for XDG_RUNTIME_DIR
MOZART_UID=$(id -u "$MOZART_USER")
sed "s|/run/user/1001|/run/user/${MOZART_UID}|" \
  "${REPO_DIR}/deploy/mozart.service" > /etc/systemd/system/mozart.service

systemctl daemon-reload
systemctl enable mozart
systemctl start mozart

# Caddy config
mkdir -p /etc/caddy
cp "${REPO_DIR}/deploy/Caddyfile" /etc/caddy/Caddyfile

# Generate a random temporary password so Caddy starts immediately
TEMP_PASSWORD=$(openssl rand -base64 16)
TEMP_HASH=$(caddy hash-password --plaintext "$TEMP_PASSWORD")

printf 'MOZART_DOMAIN=%s\nMOZART_USER=admin\nMOZART_HASH=%s\n' \
  "$CADDY_ADDRESS" "$TEMP_HASH" > /etc/caddy/env

# Point caddy's systemd unit at the env file
mkdir -p /etc/systemd/system/caddy.service.d
cat > /etc/systemd/system/caddy.service.d/env.conf <<EOF
[Service]
EnvironmentFile=/etc/caddy/env
EOF

systemctl daemon-reload
systemctl restart caddy

# Install health check cron
cp "${REPO_DIR}/deploy/health-check.sh" "${MOZART_HOME}/health-check.sh"
chmod +x "${MOZART_HOME}/health-check.sh"
CRON_LINE="*/5 * * * * ${MOZART_HOME}/health-check.sh"
(crontab -u "$MOZART_USER" -l 2>/dev/null || true; echo "$CRON_LINE") \
  | sort -u | crontab -u "$MOZART_USER" -

echo ""
echo "============================================="
echo "  Mozart is running!"
echo "============================================="
echo ""
echo "  Login credentials (change these later):"
echo "    Username: admin"
echo "    Password: ${TEMP_PASSWORD}"
echo ""
if [[ -z "$DOMAIN" ]]; then
  echo "  Access: ${CADDY_ADDRESS}"
  echo "  (HTTP only — add a domain later for TLS)"
else
  echo "  Access: https://${DOMAIN}"
  echo "  (Point DNS A record to this server first)"
fi
echo ""
echo "  To change the password:"
echo "    caddy hash-password --plaintext 'new-password'"
echo "    sudo nano /etc/caddy/env  # update MOZART_HASH"
echo "    sudo systemctl restart caddy"
echo ""
if [[ -z "$DOMAIN" ]]; then
  echo "  To upgrade to HTTPS later:"
  echo "    Edit /etc/caddy/env: set MOZART_DOMAIN=yourdomain.com"
  echo "    sudo systemctl restart caddy"
  echo ""
fi
echo "  Verify: curl -sf localhost:4141/health"
echo ""