#!/usr/bin/env bash
# G-DAIR Agent -- universal install + outbound relay setup
#
# Works on: Linux (x86_64, arm64, armv7), macOS (arm64, x86_64)
# Works as: root, sudo, or unprivileged user
# Works on: servers, laptops, VMs, containers, CI runners, offices
#
# The agent connects OUTBOUND to keys.ag every 3 seconds to receive commands.
# No inbound port is required. Works behind NAT, firewalls, Cloudflare, etc.
#
# Usage (generated by keys.ag dashboard):
#   curl -fsSL 'https://keys.ag/install.sh?t=<TOKEN>' | bash
#
# Overrides:
#   GDAIR_PORT          local agent HTTP port  (default: 7777)
#   GDAIR_INSTALL_DIR   install path           (default: auto)
#   GDAIR_BINARY_URL    binary download URL    (default: auto)

set -euo pipefail

# ── Token (injected by server) ────────────────────────────────────────────────
TOKEN="${GDAIR_TOKEN:-}"
if [[ -z "$TOKEN" ]]; then
  echo "ERROR: No install token. Generate one from https://keys.ag -> Deployments."
  exit 1
fi

# ── OS / arch / privilege detection ──────────────────────────────────────────
OS="$(uname -s 2>/dev/null || echo Linux)"
ARCH="$(uname -m 2>/dev/null || echo x86_64)"
EUID_VAL="${EUID:-$(id -u 2>/dev/null || echo 1)}"
IS_ROOT=0; [[ "$EUID_VAL" -eq 0 ]] && IS_ROOT=1

# ── Install directory -- never fail on permissions ────────────────────────────
if [[ -n "${GDAIR_INSTALL_DIR:-}" ]]; then
  INSTALL_DIR="$GDAIR_INSTALL_DIR"
elif [[ $IS_ROOT -eq 1 ]]; then
  INSTALL_DIR="/opt/gdair"
elif mkdir -p /opt/gdair 2>/dev/null && touch /opt/gdair/.w 2>/dev/null; then
  rm -f /opt/gdair/.w
  INSTALL_DIR="/opt/gdair"
elif [[ -n "${HOME:-}" && "$HOME" != "/" ]]; then
  INSTALL_DIR="$HOME/gdair"
else
  INSTALL_DIR="/tmp/gdair-$$"
fi

PORT="${GDAIR_PORT:-7777}"
REGISTRY_BASE="https://keys.ag"
BINARY_BASE="https://keys.ag/binaries"
BINARY="$INSTALL_DIR/gdair-locker"
LOG="$INSTALL_DIR/agent.log"
PID_FILE="$INSTALL_DIR/agent.pid"
HEARTBEAT_PID_FILE="$INSTALL_DIR/heartbeat.pid"
HEARTBEAT_SCRIPT="$INSTALL_DIR/heartbeat.sh"
AGENT_ID_FILE="$INSTALL_DIR/.agent-id"
AGENT_BUNDLE_FILE="$INSTALL_DIR/agent-bundle.json"
REGISTER_PAYLOAD_FILE="$INSTALL_DIR/.register-payload.json"
RELAY_CONFIG_FILE="$INSTALL_DIR/relay-config.json"

echo ""
echo "  G-DAIR Agent -- Installing"
echo "  --------------------------"
printf "  OS      : %s/%s%s\n" "$OS" "$ARCH" "$([[ $IS_ROOT -eq 1 ]] && echo ' (root)' || echo ' (user)')"
printf "  Dir     : %s\n" "$INSTALL_DIR"
printf "  Port    : %s\n" "$PORT"
echo ""

# ── Redeem install token ──────────────────────────────────────────────────────
echo "-> Redeeming install token..."
DEPLOY_JSON="$(curl -fsSL --max-time 15 "${REGISTRY_BASE}/gdair/install?t=${TOKEN}" 2>/dev/null || echo '{}')"

_json_get() {
  local key="$1" json="$2"
  python3 -c "import json,sys; print(json.loads(r'''$json''').get('$key',''))" 2>/dev/null \
    || perl -ne "if(m|\"$key\"\s*:\s*\"([^\"]+)\"|){print \$1;last}" <<<"$json" 2>/dev/null \
    || grep -o "\"$key\":\"[^\"]*\"" <<<"$json" | head -1 | sed 's/.*":"\(.*\)"/\1/' 2>/dev/null \
    || echo ""
}
_json_escape() { printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'; }

DEPLOY_OK="$(_json_get ok "$DEPLOY_JSON")"
if [[ "$DEPLOY_OK" != "True" && "$DEPLOY_OK" != "true" ]]; then
  ERR="$(_json_get error "$DEPLOY_JSON")"
  echo "  ERROR: ${ERR:-token invalid or expired}"
  echo "  Get a new URL at https://keys.ag -> Deployments -> '+ Install on new host'"
  exit 1
fi

DEPLOY_NAME="$(_json_get deployment_name "$DEPLOY_JSON")"
PUB_KEY="$(_json_get public_key_b64 "$DEPLOY_JSON")"
AGENT_SECRET="$(_json_get agent_secret_b64 "$DEPLOY_JSON")"
DEPLOY_ID="$(_json_get deployment_id "$DEPLOY_JSON")"

printf "  OK Deployment : %s\n" "$DEPLOY_NAME"
printf "  OK Public key : %.14s...\n" "$PUB_KEY"

# ── Ensure install directory ──────────────────────────────────────────────────
mkdir -p "$INSTALL_DIR" "$INSTALL_DIR/data" 2>/dev/null || {
  echo "  ERROR: cannot create $INSTALL_DIR -- set GDAIR_INSTALL_DIR to a writable path"
  exit 1
}
chmod 700 "$INSTALL_DIR" 2>/dev/null || true

# ── Stable agent ID (persists across reinstalls) ─────────────────────────────
if [[ -f "$AGENT_ID_FILE" ]]; then
  AGENT_ID="$(cat "$AGENT_ID_FILE" 2>/dev/null || echo "")"
fi
if [[ -z "${AGENT_ID:-}" ]]; then
  # Try multiple UUID sources for portability
  AGENT_ID="$(uuidgen 2>/dev/null \
    || cat /proc/sys/kernel/random/uuid 2>/dev/null \
    || python3 -c "import uuid; print(uuid.uuid4())" 2>/dev/null \
    || printf '%s-%s-%s' "$(date +%s)" "$$" "${RANDOM:-0}")"
  printf '%s\n' "$AGENT_ID" > "$AGENT_ID_FILE"
  chmod 600 "$AGENT_ID_FILE" 2>/dev/null || true
fi

# ── Platform and binary URL ───────────────────────────────────────────────────
case "${OS}-${ARCH}" in
  Linux-x86_64|Linux-amd64)  BINARY_URL="${BINARY_BASE}/gdair-locker-linux-x64-musl" ;;
  Linux-aarch64|Linux-arm64) BINARY_URL="${BINARY_BASE}/gdair-locker-linux-arm64-musl" ;;
  Linux-armv7*|Linux-armv6*) BINARY_URL="${BINARY_BASE}/gdair-locker-linux-armv7-musl" ;;
  Darwin-arm64)               BINARY_URL="${BINARY_BASE}/gdair-locker-macos-arm64" ;;
  Darwin-x86_64)              BINARY_URL="${BINARY_BASE}/gdair-locker-macos-x64" ;;
  *)                          BINARY_URL="${BINARY_BASE}/gdair-locker-linux-x64-gnu" ;;
esac
BINARY_URL="${GDAIR_BINARY_URL:-$BINARY_URL}"

# ── Download binary (always refresh) ─────────────────────────────────────────
BINARY_NEXT="$INSTALL_DIR/.gdair-locker.next"
echo "-> Downloading agent binary (${OS}/${ARCH})..."
if command -v curl >/dev/null 2>&1; then
  curl -fsSL --progress-bar "$BINARY_URL" -o "$BINARY_NEXT" 2>/dev/null \
    || curl -fsSL "$BINARY_URL" -o "$BINARY_NEXT"
elif command -v wget >/dev/null 2>&1; then
  wget -q "$BINARY_URL" -O "$BINARY_NEXT"
else
  echo "  ERROR: curl or wget required"
  exit 1
fi
chmod +x "$BINARY_NEXT"

if ! "$BINARY_NEXT" --version >/dev/null 2>&1; then
  echo "  ERROR: downloaded binary failed to execute (wrong arch?)"
  echo "  Set GDAIR_BINARY_URL to override. Expected: ${OS}/${ARCH}"
  rm -f "$BINARY_NEXT"
  exit 1
fi
AGENT_VERSION="$("$BINARY_NEXT" --version 2>/dev/null | awk '{print $NF}' || echo unknown)"
echo "  OK Downloaded v${AGENT_VERSION}"

# ── Stop existing processes ───────────────────────────────────────────────────
for pidfile in "$PID_FILE" "$HEARTBEAT_PID_FILE"; do
  if [[ -f "$pidfile" ]]; then
    OLD="$(cat "$pidfile" 2>/dev/null || echo "")"
    [[ -n "$OLD" ]] && kill -0 "$OLD" 2>/dev/null && kill "$OLD" 2>/dev/null && sleep 1 || true
    rm -f "$pidfile"
  fi
done

mv -f "$BINARY_NEXT" "$BINARY"
chmod +x "$BINARY"

# ── Write agent bundle (keys for the agent to start with) ────────────────────
# This is equivalent to a lightweight .unlock file -- no master private key.
cat > "$AGENT_BUNDLE_FILE" <<BUNDLE
{
  "version": "0.2.0",
  "scheme": "X25519+ChaCha20-Poly1305",
  "deployment_name": "$(_json_escape "$DEPLOY_NAME")",
  "deployment_id": "$(_json_escape "$DEPLOY_ID")",
  "public_key_b64": "$(_json_escape "$PUB_KEY")",
  "agent_secret_b64": "$(_json_escape "$AGENT_SECRET")"
}
BUNDLE
chmod 600 "$AGENT_BUNDLE_FILE"

# ── Write relay config (enables outbound command relay to keys.ag) ────────────
# The relay polls keys.ag every 3s and executes commands sent from the dashboard.
# This works behind NAT, corporate firewalls, and mobile networks -- no inbound needed.
cat > "$RELAY_CONFIG_FILE" <<RELAY
{
  "agent_id": "$(_json_escape "$AGENT_ID")",
  "deployment_id": "$(_json_escape "$DEPLOY_ID")",
  "public_key_b64": "$(_json_escape "$PUB_KEY")",
  "agent_secret_b64": "$(_json_escape "$AGENT_SECRET")",
  "hostname": "$(_json_escape "$(hostname -s 2>/dev/null || hostname 2>/dev/null || echo unknown)")",
  "host": "pending",
  "port": $PORT,
  "platform": "$(_json_escape "${OS}/${ARCH}")",
  "agent_version": "$(_json_escape "$AGENT_VERSION")",
  "install_type": "script",
  "transport": "outbound-relay",
  "permitted_origins": ["https://keys.ag"]
}
RELAY
chmod 600 "$RELAY_CONFIG_FILE"

# ── Start agent with outbound relay enabled ───────────────────────────────────
echo "-> Starting agent on 0.0.0.0:${PORT} with outbound relay..."
GDAIR_AGENT_ID="$AGENT_ID" nohup "$BINARY" agent start \
  --unlock-file "$AGENT_BUNDLE_FILE" \
  --bind "0.0.0.0:${PORT}" \
  --scan-dir "$INSTALL_DIR/data" \
  --registry "$INSTALL_DIR/registry.json" \
  --relay-config "$RELAY_CONFIG_FILE" \
  > "$LOG" 2>&1 &
echo $! > "$PID_FILE"
sleep 3

AGENT_PID="$(cat "$PID_FILE" 2>/dev/null || echo 0)"
if ! kill -0 "$AGENT_PID" 2>/dev/null; then
  echo "  ERROR: agent failed to start"
  tail -15 "$LOG" 2>/dev/null
  exit 1
fi
echo "  OK Agent running (PID $AGENT_PID)"

# ── Firewall (root only) ──────────────────────────────────────────────────────
if [[ $IS_ROOT -eq 1 ]]; then
  if command -v ufw >/dev/null 2>&1 && ufw status 2>/dev/null | grep -q "Status: active"; then
    ufw allow "${PORT}/tcp" >/dev/null 2>&1 && echo "  OK ufw: opened TCP ${PORT}"
  fi
  iptables -I INPUT -p tcp --dport "${PORT}" -j ACCEPT 2>/dev/null || true
else
  echo "  INFO: Not root -- skipping firewall. Agent uses outbound relay (no inbound needed)."
fi

# ── Public IP + registration payload ─────────────────────────────────────────
MY_IP="$(curl -4 -sf --max-time 5 https://ifconfig.me 2>/dev/null \
      || curl -4 -sf --max-time 5 https://api.ipify.org 2>/dev/null \
      || hostname -I 2>/dev/null | awk '{print $1}' \
      || ipconfig getifaddr en0 2>/dev/null \
      || echo "unknown")"
HOSTNAME_VALUE="$(hostname -s 2>/dev/null || hostname 2>/dev/null || echo "$AGENT_ID")"

# Update relay config with real IP now that we know it
if command -v python3 >/dev/null 2>&1; then
  python3 - <<PY 2>/dev/null || true
import json
path = "$RELAY_CONFIG_FILE"
try:
    d = json.load(open(path))
    d["host"] = "$MY_IP"
    json.dump(d, open(path, "w"), indent=2)
except: pass
PY
fi

# Build registration payload once -- reused by heartbeat
cat > "$REGISTER_PAYLOAD_FILE" <<PAYLOAD
{
  "agent_id": "$(_json_escape "$AGENT_ID")",
  "deployment_id": "$(_json_escape "$DEPLOY_ID")",
  "public_key_b64": "$(_json_escape "$PUB_KEY")",
  "agent_secret_b64": "$(_json_escape "$AGENT_SECRET")",
  "hostname": "$(_json_escape "$HOSTNAME_VALUE")",
  "host": "$(_json_escape "$MY_IP")",
  "port": $PORT,
  "platform": "$(_json_escape "${OS}/${ARCH}")",
  "agent_version": "$(_json_escape "$AGENT_VERSION")",
  "install_type": "script",
  "transport": "outbound-relay"
}
PAYLOAD
chmod 600 "$REGISTER_PAYLOAD_FILE"

# ── Register now ─────────────────────────────────────────────────────────────
echo "-> Registering with keys.ag (${MY_IP}:${PORT})..."
REG="$(curl -sf --max-time 10 \
  -X POST "${REGISTRY_BASE}/gdair/register" \
  -H "Content-Type: application/json" \
  --data-binary "@$REGISTER_PAYLOAD_FILE" \
  2>/dev/null || echo '{}')"
if echo "$REG" | grep -q '"registered":true'; then
  echo "  OK Registered -- agent appears online in the dashboard"
else
  echo "  WARN: Registration pending -- agent will appear when reachable"
fi

# ── Heartbeat: keeps registry presence alive (survives NAT timeouts) ──────────
# Uses the PID of the running agent to self-stop when the agent dies.
# Written as a separate file to avoid heredoc variable expansion bugs.
cat > "$HEARTBEAT_SCRIPT" <<'HSEOF'
#!/usr/bin/env bash
PIDFILE="__PID_FILE__"
PAYLOAD="__REGISTER_PAYLOAD_FILE__"
REGISTRY="__REGISTRY_BASE__"
while true; do
  sleep 30
  # Stop heartbeat if agent process has died
  AGENT_PID="$(cat "$PIDFILE" 2>/dev/null || echo 0)"
  if ! kill -0 "$AGENT_PID" 2>/dev/null; then
    exit 0
  fi
  curl -sf --max-time 10 -X POST "${REGISTRY}/gdair/register" \
    -H "Content-Type: application/json" \
    --data-binary "@$PAYLOAD" >/dev/null 2>&1 || true
done
HSEOF
# Replace placeholders (avoids heredoc variable expansion issues)
sed -i.bak \
  -e "s|__PID_FILE__|$PID_FILE|g" \
  -e "s|__REGISTER_PAYLOAD_FILE__|$REGISTER_PAYLOAD_FILE|g" \
  -e "s|__REGISTRY_BASE__|$REGISTRY_BASE|g" \
  "$HEARTBEAT_SCRIPT" 2>/dev/null || \
  sed \
  -e "s|__PID_FILE__|$PID_FILE|g" \
  -e "s|__REGISTER_PAYLOAD_FILE__|$REGISTER_PAYLOAD_FILE|g" \
  -e "s|__REGISTRY_BASE__|$REGISTRY_BASE|g" \
  "$HEARTBEAT_SCRIPT" > "${HEARTBEAT_SCRIPT}.new" && mv "${HEARTBEAT_SCRIPT}.new" "$HEARTBEAT_SCRIPT"
rm -f "${HEARTBEAT_SCRIPT}.bak"
chmod 700 "$HEARTBEAT_SCRIPT"

nohup bash "$HEARTBEAT_SCRIPT" > "$INSTALL_DIR/heartbeat.log" 2>&1 &
echo $! > "$HEARTBEAT_PID_FILE"
echo "  OK Heartbeat running (keeps dashboard presence alive)"

# ── Summary ───────────────────────────────────────────────────────────────────
echo ""
echo "  +------------------------------------------------------+"
echo "  |  Agent running -- outbound relay active              |"
echo "  +------------------------------------------------------+"
printf "  |  Deployment  : %s\n" "$DEPLOY_NAME"
printf "  |  Host        : %s\n" "${MY_IP}"
printf "  |  Port        : %s\n" "$PORT"
printf "  |  Agent ID    : %.36s\n" "${AGENT_ID}"
printf "  |  Install dir : %s\n" "$INSTALL_DIR"
echo "  |"
echo "  |  Commands from keys.ag reach this agent via"
echo "  |  outbound HTTPS relay -- no inbound ports needed."
echo "  |  View at: https://keys.ag -> Home (agent appears online)"
echo "  +------------------------------------------------------+"
echo ""
printf "  Logs : tail -f %s\n" "$LOG"
printf "  Stop : kill \$(cat %s)\n" "$PID_FILE"
echo ""

# macOS: suggest launchd for persistence across logout
if [[ "$OS" == "Darwin" ]]; then
  PLIST_LABEL="ag.keys.gdair.agent"
  PLIST_PATH=""
  if [[ $IS_ROOT -eq 1 ]]; then
    PLIST_PATH="/Library/LaunchDaemons/${PLIST_LABEL}.plist"
  else
    mkdir -p "$HOME/Library/LaunchAgents" 2>/dev/null || true
    PLIST_PATH="$HOME/Library/LaunchAgents/${PLIST_LABEL}.plist"
  fi
  cat > "$PLIST_PATH" <<PLIST
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
  <key>Label</key><string>${PLIST_LABEL}</string>
  <key>ProgramArguments</key>
  <array>
    <string>${BINARY}</string>
    <string>agent</string><string>start</string>
    <string>--unlock-file</string><string>${AGENT_BUNDLE_FILE}</string>
    <string>--bind</string><string>0.0.0.0:${PORT}</string>
    <string>--scan-dir</string><string>${INSTALL_DIR}/data</string>
    <string>--registry</string><string>${INSTALL_DIR}/registry.json</string>
    <string>--relay-config</string><string>${RELAY_CONFIG_FILE}</string>
  </array>
  <key>EnvironmentVariables</key>
  <dict><key>GDAIR_AGENT_ID</key><string>${AGENT_ID}</string></dict>
  <key>RunAtLoad</key><true/>
  <key>KeepAlive</key><true/>
  <key>StandardOutPath</key><string>${LOG}</string>
  <key>StandardErrorPath</key><string>${LOG}</string>
</dict>
</plist>
PLIST
  if [[ $IS_ROOT -eq 1 ]]; then
    launchctl load "$PLIST_PATH" 2>/dev/null && echo "  OK launchd: agent will restart on reboot (system daemon)" || true
  else
    launchctl load "$PLIST_PATH" 2>/dev/null && echo "  OK launchd: agent will restart on login (user agent)" || true
    echo "  macOS: For system-wide persistence (survives logout), run as root."
  fi
fi

# Linux: suggest systemd for persistence
if [[ "$OS" == "Linux" ]] && [[ $IS_ROOT -eq 1 ]] && command -v systemctl >/dev/null 2>&1; then
  UNIT_FILE="/etc/systemd/system/gdair-agent.service"
  cat > "$UNIT_FILE" <<UNIT
[Unit]
Description=G-DAIR Agent (outbound relay)
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
Restart=always
RestartSec=10
Environment=GDAIR_AGENT_ID=${AGENT_ID}
ExecStart=${BINARY} agent start \
  --unlock-file ${AGENT_BUNDLE_FILE} \
  --bind 0.0.0.0:${PORT} \
  --scan-dir ${INSTALL_DIR}/data \
  --registry ${INSTALL_DIR}/registry.json \
  --relay-config ${RELAY_CONFIG_FILE}
StandardOutput=append:${LOG}
StandardError=append:${LOG}

[Install]
WantedBy=multi-user.target
UNIT
  # Stop nohup process -- systemd takes over
  kill "$AGENT_PID" 2>/dev/null || true
  systemctl daemon-reload
  systemctl enable gdair-agent 2>/dev/null || true
  systemctl start gdair-agent
  sleep 2
  if systemctl is-active gdair-agent >/dev/null 2>&1; then
    echo "  OK systemd: gdair-agent enabled and running (survives reboots)"
  else
    echo "  WARN: systemd unit failed -- nohup process still running"
    # Restart nohup fallback
    GDAIR_AGENT_ID="$AGENT_ID" nohup "$BINARY" agent start \
      --unlock-file "$AGENT_BUNDLE_FILE" \
      --bind "0.0.0.0:${PORT}" \
      --scan-dir "$INSTALL_DIR/data" \
      --registry "$INSTALL_DIR/registry.json" \
      --relay-config "$RELAY_CONFIG_FILE" \
      >> "$LOG" 2>&1 &
    echo $! > "$PID_FILE"
  fi
fi