#!/bin/sh
# shellcheck disable=SC1091
# AmyForge Hostagent Installer
# Usage: curl -sfL https://get-amyforge.amybot.ai | sh -s -- --token TOKEN --controlplane URL

# Wrap everything in main() so sh reads the entire script before executing.
# This prevents partial-execution issues when piping curl output to sh.
main() {
set -eu

# --- Defaults ---
AMYFORGE_TOKEN=""
AMYFORGE_CONTROLPLANE=""
AMYFORGE_REGION="default"
AMYFORGE_INTERFACE=""
AMYFORGE_CIDR="172.16.0.0/16"
AMYFORGE_UNINSTALL=false

BASE_URL="https://get-amyforge.amybot.ai"
INSTALL_DIR="/usr/local/bin"
DATA_DIR="/var/lib/amyforge"
CONFIG_DIR="/etc/amyforge"
SERVICE_NAME="amyforge-hostagent"

# --- Colors ---
if [ -t 1 ] || [ -t 2 ]; then
    RED=$(printf '\033[0;31m')
    GREEN=$(printf '\033[0;32m')
    YELLOW=$(printf '\033[0;33m')
    CYAN=$(printf '\033[0;36m')
    BOLD=$(printf '\033[1m')
    NC=$(printf '\033[0m')
else
    RED="" GREEN="" YELLOW="" CYAN="" BOLD="" NC=""
fi

info()  { printf '%s[INFO]%s  %s\n' "$CYAN" "$NC" "$*"; }
warn()  { printf '%s[WARN]%s  %s\n' "$YELLOW" "$NC" "$*"; }
error() { printf '%s[ERROR]%s %s\n' "$RED" "$NC" "$*"; }
ok()    { printf '%s[OK]%s    %s\n' "$GREEN" "$NC" "$*"; }

fatal() {
    error "$*"
    exit 1
}

# --- Parse arguments ---
while [ $# -gt 0 ]; do
    case "$1" in
        --token)       AMYFORGE_TOKEN="$2";       shift 2 ;;
        --controlplane) AMYFORGE_CONTROLPLANE="$2"; shift 2 ;;
        --region)      AMYFORGE_REGION="$2";       shift 2 ;;
        --interface)   AMYFORGE_INTERFACE="$2";    shift 2 ;;
        --cidr)        AMYFORGE_CIDR="$2";         shift 2 ;;
        --uninstall)   AMYFORGE_UNINSTALL=true;    shift   ;;
        --help|-h)
            cat <<'USAGE'
AmyForge Hostagent Installer

Usage:
  curl -sfL https://get-amyforge.amybot.ai | sh -s -- [OPTIONS]

Required:
  --token TOKEN            Control plane auth token
  --controlplane URL       Control plane URL (e.g. https://forge.amybot.ai)

Optional:
  --region REGION          Region identifier (default: "default")
  --interface IFACE        External network interface (auto-detected if omitted)
  --cidr CIDR              VM network CIDR (default: 172.16.0.0/16)
  --uninstall              Remove hostagent, firecracker, systemd service, and dirs

Examples:
  # Install
  curl -sfL https://get-amyforge.amybot.ai | sh -s -- \
    --token YOUR_TOKEN \
    --controlplane https://forge.amybot.ai

  # Uninstall
  curl -sfL https://get-amyforge.amybot.ai | sh -s -- --uninstall
USAGE
            exit 0
            ;;
        *)
            fatal "Unknown option: $1 (use --help for usage)"
            ;;
    esac
done

# --- Uninstall ---
do_uninstall() {
    info "Uninstalling AmyForge hostagent..."

    # Stop and disable service
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        info "Stopping $SERVICE_NAME..."
        systemctl stop "$SERVICE_NAME"
    fi
    if systemctl is-enabled --quiet "$SERVICE_NAME" 2>/dev/null; then
        info "Disabling $SERVICE_NAME..."
        systemctl disable "$SERVICE_NAME"
    fi

    # Remove service file
    if [ -f "/etc/systemd/system/${SERVICE_NAME}.service" ]; then
        rm -f "/etc/systemd/system/${SERVICE_NAME}.service"
        systemctl daemon-reload
        ok "Removed systemd service"
    fi

    # Remove config
    if [ -d "$CONFIG_DIR" ]; then
        rm -rf "$CONFIG_DIR"
        ok "Removed $CONFIG_DIR"
    fi

    # Remove sysctl config
    if [ -f "/etc/sysctl.d/99-amyforge.conf" ]; then
        rm -f "/etc/sysctl.d/99-amyforge.conf"
        ok "Removed sysctl config"
    fi

    # Remove binaries
    for bin in hostagent firecracker jailer; do
        if [ -f "${INSTALL_DIR}/${bin}" ]; then
            rm -f "${INSTALL_DIR}/${bin}"
            ok "Removed ${INSTALL_DIR}/${bin}"
        fi
    done

    # Ask about data directory
    if [ -d "$DATA_DIR" ]; then
        printf '%sRemove %s? This contains VM data (images, overlays, volumes). [y/N] %s' "$YELLOW" "$DATA_DIR" "$NC"
        read -r answer
        case "$answer" in
            [yY]|[yY][eE][sS])
                rm -rf "$DATA_DIR"
                ok "Removed $DATA_DIR"
                ;;
            *)
                info "Kept $DATA_DIR"
                ;;
        esac
    fi

    ok "AmyForge hostagent uninstalled"
    exit 0
}

if [ "$AMYFORGE_UNINSTALL" = true ]; then
    if [ "$(id -u)" -ne 0 ]; then
        fatal "Uninstall must be run as root"
    fi
    do_uninstall
fi

# --- Validate required args ---
if [ -z "$AMYFORGE_TOKEN" ]; then
    fatal "Missing required flag: --token (use --help for usage)"
fi
if [ -z "$AMYFORGE_CONTROLPLANE" ]; then
    fatal "Missing required flag: --controlplane (use --help for usage)"
fi

# --- Pre-flight checks ---
printf '\n%sAmyForge Hostagent Installer%s\n\n' "$BOLD" "$NC"

# Root check
if [ "$(id -u)" -ne 0 ]; then
    fatal "This script must be run as root (try: sudo sh -s ...)"
fi

# Detect OS
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian|pop|linuxmint)  OS_FAMILY="debian" ;;
            rhel|centos|fedora|rocky|alma|ol) OS_FAMILY="rhel" ;;
            arch|manjaro|endeavouros)     OS_FAMILY="arch" ;;
            *)
                if [ -n "${ID_LIKE:-}" ]; then
                    case "$ID_LIKE" in
                        *debian*|*ubuntu*) OS_FAMILY="debian" ;;
                        *rhel*|*fedora*|*centos*) OS_FAMILY="rhel" ;;
                        *arch*) OS_FAMILY="arch" ;;
                        *) fatal "Unsupported OS: $ID (ID_LIKE=$ID_LIKE)" ;;
                    esac
                else
                    fatal "Unsupported OS: $ID"
                fi
                ;;
        esac
        OS_NAME="$PRETTY_NAME"
    else
        fatal "Cannot detect OS: /etc/os-release not found"
    fi
}

detect_os
ok "OS: $OS_NAME ($OS_FAMILY)"

# Detect architecture
ARCH=$(uname -m)
case "$ARCH" in
    x86_64)  ARCH="x86_64" ;;
    aarch64) ARCH="aarch64" ;;
    *)       fatal "Unsupported architecture: $ARCH (need x86_64 or aarch64)" ;;
esac
ok "Architecture: $ARCH"

# Verify KVM
if [ ! -e /dev/kvm ]; then
    fatal "/dev/kvm not found. This host does not support KVM virtualization.
    Ensure you are running on bare metal or a VM with nested virtualization enabled."
fi
if [ ! -r /dev/kvm ] || [ ! -w /dev/kvm ]; then
    fatal "/dev/kvm is not accessible. Check permissions."
fi
ok "KVM: available"

# Verify cgroups v2
if [ ! -f /sys/fs/cgroup/cgroup.controllers ]; then
    fatal "cgroups v2 not available. AmyForge requires cgroups v2.
    Check your kernel boot parameters (systemd.unified_cgroup_hierarchy=1)."
fi
ok "Cgroups v2: available"

# --- Install system dependencies ---
info "Installing system dependencies..."

install_deps_debian() {
    apt-get update -qq >/dev/null 2>&1
    apt-get install -y -qq iptables iproute2 ca-certificates curl >/dev/null 2>&1
}

install_deps_rhel() {
    dnf install -y -q iptables iproute ca-certificates curl >/dev/null 2>&1
}

install_deps_arch() {
    pacman -Sy --noconfirm --needed iptables iproute2 ca-certificates curl >/dev/null 2>&1
}

case "$OS_FAMILY" in
    debian) install_deps_debian ;;
    rhel)   install_deps_rhel ;;
    arch)   install_deps_arch ;;
esac
ok "System dependencies installed"

# --- Auto-detect host resources ---
HOST_CPUS=$(nproc)
HOST_MEMORY_KB=$(awk '/^MemTotal:/ { print $2 }' /proc/meminfo)
HOST_MEMORY_MB=$((HOST_MEMORY_KB / 1024))
HOST_HOSTNAME=$(hostname)
HOST_IP=""

# Detect external interface and IP
if [ -n "$AMYFORGE_INTERFACE" ]; then
    EXT_IFACE="$AMYFORGE_INTERFACE"
else
    EXT_IFACE=$(ip route show default | awk '/default/ { print $5; exit }')
    if [ -z "$EXT_IFACE" ]; then
        fatal "Could not auto-detect external interface. Use --interface to specify."
    fi
fi

HOST_IP=$(ip -4 addr show "$EXT_IFACE" | awk '/inet / { split($2, a, "/"); print a[1]; exit }')
if [ -z "$HOST_IP" ]; then
    fatal "Could not detect IP address on interface $EXT_IFACE"
fi

ok "Hostname: $HOST_HOSTNAME"
ok "CPUs: $HOST_CPUS"
ok "Memory: ${HOST_MEMORY_MB} MB"
ok "Interface: $EXT_IFACE ($HOST_IP)"

# --- Download binaries ---
info "Downloading binaries for $ARCH..."

BIN_URL="${BASE_URL}/bin/${ARCH}"

download() {
    _dl_url="$1"
    _dl_dest="$2"
    _dl_name=$(basename "$_dl_dest")

    if ! curl -sfL "$_dl_url" -o "$_dl_dest"; then
        fatal "Failed to download $_dl_name from $_dl_url"
    fi
}

download "${BIN_URL}/hostagent"   "${INSTALL_DIR}/hostagent"
ok "Downloaded hostagent"
download "${BIN_URL}/firecracker" "${INSTALL_DIR}/firecracker"
ok "Downloaded firecracker"
download "${BIN_URL}/jailer"      "${INSTALL_DIR}/jailer"
ok "Downloaded jailer"

chmod +x "${INSTALL_DIR}/hostagent" "${INSTALL_DIR}/firecracker" "${INSTALL_DIR}/jailer"

# --- Download kernel ---
info "Downloading kernel image..."
mkdir -p "$DATA_DIR"
download "${BIN_URL}/vmlinux" "${DATA_DIR}/vmlinux"
ok "Downloaded vmlinux kernel"

# --- Create directories ---
mkdir -p "${DATA_DIR}/images" "${DATA_DIR}/overlays" "${DATA_DIR}/volumes" "${DATA_DIR}/snapshots"

# --- Download base rootfs image ---
info "Downloading base rootfs image..."
download "${BASE_URL}/bin/${ARCH}/images/alpine-latest.ext4" "${DATA_DIR}/images/alpine-latest.ext4"
ok "Downloaded base rootfs image (alpine-latest.ext4)"

# --- Download golden snapshot ---
info "Downloading OpenClaw golden snapshot (this may take a few minutes)..."
SNAP_URL="${BIN_URL}/snapshots/openclaw-golden"
SNAP_DIR="${DATA_DIR}/snapshots/openclaw-golden"
mkdir -p "$SNAP_DIR"
download "${SNAP_URL}/rootfs.ext4"          "${SNAP_DIR}/rootfs.ext4"
ok "Downloaded rootfs.ext4"
download "${SNAP_URL}/memory.snap"          "${SNAP_DIR}/memory.snap"
ok "Downloaded memory.snap"
download "${SNAP_URL}/vmstate.snap"         "${SNAP_DIR}/vmstate.snap"
ok "Downloaded vmstate.snap"
download "${SNAP_URL}/snapshot_meta.json"   "${SNAP_DIR}/snapshot_meta.json"
ok "Downloaded snapshot_meta.json"
ok "Golden snapshot ready"

mkdir -p "$CONFIG_DIR"
ok "Created directories"

# --- Write config ---
info "Writing config to ${CONFIG_DIR}/hostagent.yaml..."

cat > "${CONFIG_DIR}/hostagent.yaml" <<EOF
controlplane:
  address: ${AMYFORGE_CONTROLPLANE}
agent:
  address: "${HOST_IP}"
  grpc_port: 9091
  region: ${AMYFORGE_REGION}
  token: "${AMYFORGE_TOKEN}"
firecracker:
  binary: ${INSTALL_DIR}/firecracker
  jailer: ${INSTALL_DIR}/jailer
  kernel: ${DATA_DIR}/vmlinux
storage:
  images_dir: ${DATA_DIR}/images
  overlays_dir: ${DATA_DIR}/overlays
  volumes_dir: ${DATA_DIR}/volumes
  snapshots_dir: ${DATA_DIR}/snapshots
network:
  cidr: ${AMYFORGE_CIDR}
  external_interface: ${EXT_IFACE}
log:
  level: info
  format: json
EOF
ok "Config written"

# --- Enable IP forwarding ---
info "Enabling IP forwarding..."
sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1
cat > /etc/sysctl.d/99-amyforge.conf <<EOF
net.ipv4.ip_forward=1
EOF
ok "IP forwarding enabled"

# --- Create systemd service ---
info "Creating systemd service..."

cat > "/etc/systemd/system/${SERVICE_NAME}.service" <<EOF
[Unit]
Description=AmyForge Host Agent
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=${INSTALL_DIR}/hostagent --config ${CONFIG_DIR}/hostagent.yaml
Restart=always
RestartSec=5
LimitNOFILE=65536
LimitNPROC=65536

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
ok "Systemd service created"

# --- Start service ---
info "Starting $SERVICE_NAME..."
systemctl enable --now "$SERVICE_NAME"
ok "Service started"

# --- Register with control plane ---
info "Registering host with control plane..."

# Strip trailing slash from controlplane URL
CP_URL=$(echo "$AMYFORGE_CONTROLPLANE" | sed 's:/*$::')

REGISTER_BODY=$(cat <<EOF
{
  "hostname": "${HOST_HOSTNAME}",
  "address": "${HOST_IP}",
  "grpc_port": 9091,
  "total_cpus": ${HOST_CPUS},
  "total_memory_mb": ${HOST_MEMORY_MB},
  "cidr_range": "${AMYFORGE_CIDR}",
  "region": "${AMYFORGE_REGION}"
}
EOF
)

REGISTER_RESP=$(curl -sf -X POST \
    "${CP_URL}/api/v1/hosts/register" \
    -H "Authorization: Bearer ${AMYFORGE_TOKEN}" \
    -H "Content-Type: application/json" \
    -d "$REGISTER_BODY" 2>&1) || {
    warn "Failed to register with control plane. The service is running and will"
    warn "connect via gRPC heartbeat. You can register manually later."
    warn "Response: $REGISTER_RESP"
}

HOST_ID=""
if [ -n "$REGISTER_RESP" ]; then
    # Extract host ID from JSON response (portable — no jq dependency)
    HOST_ID=$(echo "$REGISTER_RESP" | sed -n 's/.*"id"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')
fi

# --- Register base image with control plane ---
info "Registering base image..."
IMAGE_SIZE=$(stat -c%s "${DATA_DIR}/images/alpine-latest.ext4" 2>/dev/null || echo "104857600")
IMAGE_RESP=$(curl -sfL -X POST "${CP_URL}/api/v1/images" \
    -H "Authorization: Bearer ${AMYFORGE_TOKEN}" \
    -H "Content-Type: application/json" \
    -d "{
      \"name\": \"alpine\",
      \"version\": \"latest\",
      \"path\": \"/var/lib/amyforge/images/alpine-latest.ext4\",
      \"size_bytes\": ${IMAGE_SIZE},
      \"os\": \"linux\",
      \"arch\": \"${ARCH}\"
    }" 2>/dev/null || true)

if echo "$IMAGE_RESP" | grep -q '"id"'; then
    IMAGE_ID=$(echo "$IMAGE_RESP" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    ok "Base image registered: ${IMAGE_ID}"
else
    # Image may already exist — try to fetch existing image ID
    info "Image may already exist, fetching existing image list..."
    IMAGES_RESP=$(curl -sfL "${CP_URL}/api/v1/images" \
        -H "Authorization: Bearer ${AMYFORGE_TOKEN}" 2>/dev/null || true)
    if [ -n "$IMAGES_RESP" ]; then
        IMAGE_ID=$(echo "$IMAGES_RESP" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
        if [ -n "$IMAGE_ID" ]; then
            ok "Found existing image: ${IMAGE_ID}"
        fi
    fi
fi

# Create symlink so hostagent can find image by UUID
# The hostagent looks for {image_id}.ext4 but we downloaded alpine-latest.ext4
if [ -n "${IMAGE_ID:-}" ]; then
    ln -sf "alpine-latest.ext4" "${DATA_DIR}/images/${IMAGE_ID}.ext4"
    ok "Created image symlink: ${IMAGE_ID}.ext4 -> alpine-latest.ext4"
else
    warn "Could not determine image ID — you may need to manually symlink the image file"
fi

# --- Register OpenClaw image (snapshot-based) ---
info "Registering OpenClaw image..."
OC_IMAGE_RESP=$(curl -sfL -X POST "${CP_URL}/api/v1/images" \
    -H "Authorization: Bearer ${AMYFORGE_TOKEN}" \
    -H "Content-Type: application/json" \
    -d "{
      \"name\": \"openclaw\",
      \"version\": \"latest\",
      \"path\": \"${DATA_DIR}/snapshots/openclaw-golden/rootfs.ext4\",
      \"size_bytes\": $(stat -c%s "${SNAP_DIR}/rootfs.ext4" 2>/dev/null || echo "4294967296"),
      \"os\": \"linux\",
      \"arch\": \"${ARCH}\",
      \"snapshot_enabled\": true,
      \"snapshot_id\": \"openclaw-golden\"
    }" 2>/dev/null || true)

if echo "$OC_IMAGE_RESP" | grep -q '"id"'; then
    OC_IMAGE_ID=$(echo "$OC_IMAGE_RESP" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    ok "OpenClaw image registered: ${OC_IMAGE_ID}"
else
    info "OpenClaw image may already exist"
fi

# --- Print summary ---
printf '\n%s%sAmyForge hostagent installed successfully!%s\n\n' "$BOLD" "$GREEN" "$NC"

printf '%sInstallation:%s\n' "$BOLD" "$NC"
printf "  Hostagent:    %s/hostagent\n" "$INSTALL_DIR"
printf "  Firecracker:  %s/firecracker\n" "$INSTALL_DIR"
printf "  Jailer:       %s/jailer\n" "$INSTALL_DIR"
printf "  Kernel:       %s/vmlinux\n" "$DATA_DIR"
printf "  Config:       %s/hostagent.yaml\n" "$CONFIG_DIR"

printf '\n%sHost Resources:%s\n' "$BOLD" "$NC"
printf "  Hostname:     %s\n" "$HOST_HOSTNAME"
printf "  IP:           %s\n" "$HOST_IP"
printf "  CPUs:         %s\n" "$HOST_CPUS"
printf "  Memory:       %s MB\n" "$HOST_MEMORY_MB"
printf "  Interface:    %s\n" "$EXT_IFACE"
printf "  VM CIDR:      %s\n" "$AMYFORGE_CIDR"
printf "  Region:       %s\n" "$AMYFORGE_REGION"

if [ -n "$HOST_ID" ]; then
    printf '\n%sControl Plane:%s\n' "$BOLD" "$NC"
    printf "  URL:          %s\n" "$CP_URL"
    printf "  Host ID:      %s\n" "$HOST_ID"
fi

printf '\n%sManagement:%s\n' "$BOLD" "$NC"
printf "  Status:       systemctl status %s\n" "$SERVICE_NAME"
printf "  Logs:         journalctl -u %s -f\n" "$SERVICE_NAME"
printf "  Restart:      systemctl restart %s\n" "$SERVICE_NAME"
printf "  Uninstall:    curl -sfL %s | sh -s -- --uninstall\n" "$BASE_URL"
printf "\n"
}

# Run main with all script arguments
main "$@"