#!/bin/bash
set -euo pipefail

# --- Constants ---
NSPAWN_CONF_DIR="/etc/systemd/nspawn"
PERSISTENT_NSPAWN_DIR="/data/custom/nspawn"
PERSISTENT_DPKG_DIR="/data/custom/dpkg"
PERSISTENT_META_DIR="/data/custom/twingate"
BOOT_SCRIPT_DIR="/data/on_boot.d"
BOOT_TIMEOUT=300

# Note: This can and should be changed inside the container after setup
CONTAINER_PASSWORD="twingate"

# --- Root check ---
if [[ $EUID -ne 0 ]]; then
    echo "Error: This script must be run as root." >&2
    exit 1
fi

# --- Cleanup trap ---
SETUP_STARTED=false
cleanup() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]] && [[ "$SETUP_STARTED" == "true" ]]; then
        echo ""
        echo "Error: Setup failed (exit code $exit_code)."
        echo "If the container was partially created, you can clean up with:"
        echo "  rm -rf ${MACHINE_DIR:-/data/custom/machines/UNKNOWN}"
        echo "  rm -f ${MACHINES_LINK:-} ${NSPAWN_CONF:-}"
        echo "Then re-run setup.sh."
    fi
}
trap cleanup EXIT

# --- Helper: run a command inside the running container via nsenter ---
container_exec() {
    local pid
    pid=$(machinectl show "$CONTAINER_NAME" -p Leader --value)
    nsenter -t "$pid" -m -u -i -n -p -- "$@"
}

# --- Helper: install and enable the unifi-common boot service ---
ensure_boot_service() {
    if systemctl list-unit-files udm-boot.service &>/dev/null \
       && systemctl is-enabled udm-boot.service &>/dev/null; then
        echo "==> udm-boot.service already installed and enabled."
        return 0
    fi
    echo "==> Installing unifi-common boot service..."
    curl -fsL "https://raw.githubusercontent.com/unifi-utilities/unifi-common/HEAD/remote_install.sh" | /bin/bash
    if ! systemctl is-enabled udm-boot.service &>/dev/null; then
        echo "Warning: udm-boot.service could not be installed."
        echo "Boot persistence will NOT work. The container will not auto-restore after firmware updates."
        echo "Install manually: https://github.com/unifi-utilities/unifi-common"
        return 1
    fi
}

# --- Helper: cache host packages for offline boot recovery ---
cache_host_packages() {
    echo "==> Caching host packages for offline boot recovery..."
    mkdir -p "$PERSISTENT_DPKG_DIR"
    cp /var/cache/apt/archives/*.deb "$PERSISTENT_DPKG_DIR/" 2>/dev/null || true
    echo "  Cached $(find "$PERSISTENT_DPKG_DIR" -name '*.deb' 2>/dev/null | wc -l | tr -d ' ') packages."
}

# --- Helper: persist the nspawn config ---
persist_nspawn_config() {
    echo "==> Persisting nspawn config to ${PERSISTENT_NSPAWN_CONF}..."
    mkdir -p "$PERSISTENT_NSPAWN_DIR"
    cp "$NSPAWN_CONF" "$PERSISTENT_NSPAWN_CONF"
}

# --- Helper: save container metadata ---
save_metadata() {
    echo "==> Saving container metadata..."
    mkdir -p "$PERSISTENT_META_DIR"
    cat > "${PERSISTENT_META_DIR}/${CONTAINER_NAME}.conf" <<EOF
CONTAINER_NAME=${CONTAINER_NAME}
MACHINE_DIR=${MACHINE_DIR}
MACHINES_LINK=${MACHINES_LINK}
NSPAWN_CONF=${NSPAWN_CONF}
PERSISTENT_NSPAWN_CONF=${PERSISTENT_NSPAWN_CONF}
BOOT_SCRIPT=${BOOT_SCRIPT}
EOF
}

# --- Helper: generate the boot script ---
generate_boot_script() {
    echo "==> Generating boot script at ${BOOT_SCRIPT}..."
    mkdir -p "$BOOT_SCRIPT_DIR"

    # Variables section (expanded now)
    cat > "$BOOT_SCRIPT" <<BOOTVARS
#!/bin/bash
# Twingate Connector boot restoration script
# Generated by setup.sh -- restores ephemeral state after UniFi OS firmware updates.

CONTAINER_NAME="${CONTAINER_NAME}"
MACHINE_DIR="${MACHINE_DIR}"
MACHINES_LINK="${MACHINES_LINK}"
NSPAWN_CONF="${NSPAWN_CONF}"
NSPAWN_CONF_DIR="${NSPAWN_CONF_DIR}"
PERSISTENT_NSPAWN_CONF="${PERSISTENT_NSPAWN_CONF}"
PERSISTENT_DPKG_DIR="${PERSISTENT_DPKG_DIR}"
LOG_FILE="${PERSISTENT_META_DIR}/boot.log"
BOOT_TIMEOUT=${BOOT_TIMEOUT}
BOOTVARS

    # Script body (not expanded)
    cat >> "$BOOT_SCRIPT" <<'BOOTBODY'

log() { echo "$(date '+%Y-%m-%d %H:%M:%S') [twingate-boot] $*" | tee -a "$LOG_FILE"; }

# Rotate log if over 1MB
if [[ -f "$LOG_FILE" ]] && [[ $(stat -c%s "$LOG_FILE" 2>/dev/null || echo 0) -gt 1048576 ]]; then
    mv "$LOG_FILE" "${LOG_FILE}.old"
fi

mkdir -p "$(dirname "$LOG_FILE")"
log "=== Boot script starting ==="

# Step 1: Ensure host packages are installed
if ! command -v machinectl &>/dev/null; then
    log "machinectl not found, installing systemd-container..."
    if ls "$PERSISTENT_DPKG_DIR"/*.deb &>/dev/null; then
        log "Installing from cached packages..."
        dpkg -i "$PERSISTENT_DPKG_DIR"/*.deb 2>>"$LOG_FILE" || true
        apt-get install -f -y -qq 2>>"$LOG_FILE" || true
    fi
    if ! command -v machinectl &>/dev/null; then
        log "Cached install insufficient, attempting apt-get..."
        apt-get update -qq 2>>"$LOG_FILE" && \
        apt-get install -y -qq systemd-container 2>>"$LOG_FILE" || true
    fi
    if ! command -v machinectl &>/dev/null; then
        log "FATAL: Cannot install systemd-container. Container will not start."
        exit 1
    fi
fi

# Step 2: Restore symlink
if [[ ! -e "$MACHINES_LINK" ]]; then
    log "Restoring symlink: ${MACHINES_LINK} -> ${MACHINE_DIR}"
    mkdir -p "$(dirname "$MACHINES_LINK")"
    ln -sf "$MACHINE_DIR" "$MACHINES_LINK"
else
    log "Symlink already exists: ${MACHINES_LINK}"
fi

# Step 3: Restore nspawn config
if [[ ! -f "$NSPAWN_CONF" ]]; then
    if [[ -f "$PERSISTENT_NSPAWN_CONF" ]]; then
        log "Restoring nspawn config from ${PERSISTENT_NSPAWN_CONF}"
        mkdir -p "$NSPAWN_CONF_DIR"
        cp "$PERSISTENT_NSPAWN_CONF" "$NSPAWN_CONF"
    else
        log "WARNING: No persistent nspawn config found at ${PERSISTENT_NSPAWN_CONF}"
    fi
else
    log "nspawn config already exists: ${NSPAWN_CONF}"
fi

# Step 4: Start container if not already running
if machinectl list --no-legend 2>/dev/null | grep -q "^${CONTAINER_NAME} "; then
    log "Container '${CONTAINER_NAME}' is already running."
else
    if [[ ! -d "$MACHINE_DIR" ]]; then
        log "FATAL: Container filesystem not found at ${MACHINE_DIR}. Cannot start."
        exit 1
    fi
    log "Starting container '${CONTAINER_NAME}'..."
    machinectl start "$CONTAINER_NAME" 2>>"$LOG_FILE"
    for i in $(seq 1 "$BOOT_TIMEOUT"); do
        pid=$(machinectl show "$CONTAINER_NAME" -p Leader --value 2>/dev/null) || true
        if [[ -n "$pid" ]] && nsenter -t "$pid" -m -u -i -n -p -- /bin/true 2>/dev/null; then
            break
        fi
        if [[ $i -eq "$BOOT_TIMEOUT" ]]; then
            log "WARNING: Container did not fully boot within ${BOOT_TIMEOUT} seconds."
        fi
        sleep 1
    done
fi

# Step 5: Re-enable auto-start
machinectl enable "$CONTAINER_NAME" 2>>"$LOG_FILE" || true

# Step 6: Health check
sleep 5
pid=$(machinectl show "$CONTAINER_NAME" -p Leader --value 2>/dev/null) || true
if [[ -n "$pid" ]]; then
    if nsenter -t "$pid" -m -u -i -n -p -- systemctl is-active twingate-connector &>/dev/null; then
        log "Twingate connector is running."
    else
        log "WARNING: twingate-connector service is not active. Attempting to start..."
        nsenter -t "$pid" -m -u -i -n -p -- systemctl start twingate-connector 2>>"$LOG_FILE" || true
        sleep 3
        if nsenter -t "$pid" -m -u -i -n -p -- systemctl is-active twingate-connector &>/dev/null; then
            log "Twingate connector started successfully."
        else
            log "WARNING: Could not start twingate-connector. Check container logs."
        fi
    fi
else
    log "WARNING: Could not determine container PID for health check."
fi

log "=== Boot script complete ==="
BOOTBODY

    chmod +x "$BOOT_SCRIPT"
}

# --- Helper: configure unattended-upgrades inside the container ---
configure_unattended_upgrades() {
    echo "==> Configuring automatic updates inside container..."
    container_exec /bin/bash -c "\
        export DEBIAN_FRONTEND=noninteractive && \
        apt-get update -qq && \
        apt-get install -y -qq unattended-upgrades"

    container_exec /bin/bash -c "\
        cat > /etc/apt/apt.conf.d/50unattended-upgrades <<'UUCFG'
Unattended-Upgrade::Allowed-Origins {
    \"\${distro_id}:\${distro_codename}\";
    \"\${distro_id}:\${distro_codename}-security\";
    \"\${distro_id}ESMApps:\${distro_codename}-apps-security\";
    \"\${distro_id}ESM:\${distro_codename}-infra-security\";
    \"Twingate, Inc.:stable\";
};
Unattended-Upgrade::AutoFixInterruptedDpkg \"true\";
Unattended-Upgrade::Remove-Unused-Dependencies \"true\";
UUCFG"

    container_exec systemctl enable apt-daily.timer apt-daily-upgrade.timer
    container_exec systemctl start apt-daily.timer apt-daily-upgrade.timer
}

# --- Helper: prompt for auto-update opt-in ---
should_enable_auto_update() {
    if [[ "${AUTO_UPDATE:-}" == "1" ]]; then
        return 0
    elif [[ "${AUTO_UPDATE:-}" == "0" ]]; then
        return 1
    elif [[ -t 0 ]]; then
        read -rp "Enable automatic updates for the container OS and Twingate connector? [y/N]: " auto_update_confirm
        [[ "${auto_update_confirm,,}" == "y" ]]
    else
        return 1
    fi
}

# --- Helper: print summary for retrofit mode ---
print_retrofit_summary() {
    echo ""
    echo "============================================"
    echo "  Boot persistence added successfully!"
    echo "============================================"
    echo ""
    echo "Container name:    ${CONTAINER_NAME}"
    echo "Boot script:       ${BOOT_SCRIPT}"
    echo "Persistent config: ${PERSISTENT_NSPAWN_CONF}"
    echo "Package cache:     ${PERSISTENT_DPKG_DIR}"
    echo "Boot log:          ${PERSISTENT_META_DIR}/boot.log"
    echo ""
    echo "This container will now automatically restore after firmware updates."
    if [[ "${AUTO_UPDATE_ENABLED:-false}" == "true" ]]; then
        echo "Automatic security and connector updates are enabled."
    else
        echo "Automatic updates are not enabled. Use update.sh to update manually."
    fi
    echo ""
}

# --- Credential gathering (before any system changes) ---
get_credential() {
    local var_name="$1"
    local prompt_text="$2"
    local value="${!var_name:-}"
    if [[ -z "$value" ]]; then
        if [[ -t 0 ]]; then
            read -rp "$prompt_text: " value
        fi
        if [[ -z "$value" ]]; then
            echo "Error: ${var_name} is required. Set it as an environment variable or run the script interactively." >&2
            exit 1
        fi
    fi
    echo "$value"
}

# --- Container name ---
DEFAULT_CONTAINER_NAME="twingate-connector"
if [[ -n "${CONTAINER_NAME:-}" ]]; then
    true
elif [[ -t 0 ]]; then
    echo ""
    echo "Existing Twingate containers on this host:"
    machinectl list --no-legend 2>/dev/null | grep "twingate" || echo "  (none)"
    echo ""
    read -rp "Enter a container name [${DEFAULT_CONTAINER_NAME}]: " CONTAINER_NAME
    CONTAINER_NAME="${CONTAINER_NAME:-$DEFAULT_CONTAINER_NAME}"
else
    CONTAINER_NAME="$DEFAULT_CONTAINER_NAME"
fi

if [[ ! "$CONTAINER_NAME" =~ ^[a-zA-Z0-9_-]+$ ]]; then
    echo "Error: Container name may only contain letters, numbers, hyphens, and underscores." >&2
    exit 1
fi

# Derive paths from the chosen name
MACHINE_DIR="/data/custom/machines/${CONTAINER_NAME}"
MACHINES_LINK="/var/lib/machines/${CONTAINER_NAME}"
NSPAWN_CONF="${NSPAWN_CONF_DIR}/${CONTAINER_NAME}.nspawn"
PERSISTENT_NSPAWN_CONF="${PERSISTENT_NSPAWN_DIR}/${CONTAINER_NAME}.nspawn"
BOOT_SCRIPT="${BOOT_SCRIPT_DIR}/05-nspawn-${CONTAINER_NAME}.sh"

echo "==> Container name: ${CONTAINER_NAME}"

# --- Pre-flight checks and retrofit detection ---
if machinectl list --no-legend 2>/dev/null | grep -q "^${CONTAINER_NAME} "; then
    if [[ -f "$BOOT_SCRIPT" ]]; then
        echo "Container '${CONTAINER_NAME}' is already running and boot persistence is configured."
        echo "Nothing to do. To reinstall from scratch, run uninstall.sh first."
        exit 0
    fi

    echo "==> Container '${CONTAINER_NAME}' is running but lacks boot persistence. Retrofitting..."
    ensure_boot_service || true

    echo "==> Installing host dependencies..."
    apt-get update -qq
    apt-get install -y -qq systemd-container debootstrap
    cache_host_packages

    if [[ -f "$NSPAWN_CONF" ]]; then
        persist_nspawn_config
    else
        echo "Warning: No nspawn config found at ${NSPAWN_CONF}. Boot recovery may not work correctly."
    fi

    AUTO_UPDATE_ENABLED=false
    if should_enable_auto_update; then
        configure_unattended_upgrades
        AUTO_UPDATE_ENABLED=true
    fi
    save_metadata
    generate_boot_script
    print_retrofit_summary
    exit 0
fi

if [[ -d "$MACHINE_DIR" ]]; then
    echo "Warning: Directory ${MACHINE_DIR} already exists but container is not running."
    if [[ "${FORCE:-}" == "1" ]]; then
        echo "FORCE=1 set: removing stale directory and starting fresh."
    elif [[ -t 0 ]]; then
        read -rp "Remove it and start fresh? [y/N]: " confirm
        if [[ "${confirm,,}" != "y" ]]; then
            echo "Aborted."
            exit 1
        fi
    else
        echo "Non-interactive mode: removing stale directory and starting fresh."
    fi
    rm -rf "$MACHINE_DIR"
    rm -f "$MACHINES_LINK"
    rm -f "$NSPAWN_CONF"
    rm -f "$PERSISTENT_NSPAWN_CONF"
    rm -f "$BOOT_SCRIPT"
fi

# --- Fresh install begins here ---
SETUP_STARTED=true

TWINGATE_NETWORK=$(get_credential "TWINGATE_NETWORK" "Enter your Twingate network name (e.g., mycompany)")
TWINGATE_ACCESS_TOKEN=$(get_credential "TWINGATE_ACCESS_TOKEN" "Enter your Twingate access token")
TWINGATE_REFRESH_TOKEN=$(get_credential "TWINGATE_REFRESH_TOKEN" "Enter your Twingate refresh token")

# --- Install boot service ---
ensure_boot_service || true

# --- Install host dependencies ---
echo "==> Installing host dependencies..."
apt-get update -qq
apt-get install -y -qq systemd-container debootstrap

cache_host_packages

# --- Bootstrap container ---
echo "==> Bootstrapping Debian Bookworm container (this may take ~10 minutes)..."
mkdir -p "$(dirname "$MACHINE_DIR")"
DEBIAN_FRONTEND=noninteractive debootstrap --include=systemd,dbus,curl,ca-certificates bookworm "$MACHINE_DIR"

# --- Pre-boot container configuration ---
echo "==> Configuring container..."
systemd-nspawn -D "$MACHINE_DIR" bash -c "echo 'root:${CONTAINER_PASSWORD}' | chpasswd"
systemd-nspawn -D "$MACHINE_DIR" systemctl enable systemd-networkd

cat > "${MACHINE_DIR}/etc/resolv.conf" <<EOF
nameserver 1.1.1.1
nameserver 8.8.8.8
EOF

echo "$CONTAINER_NAME" > "${MACHINE_DIR}/etc/hostname"

# --- Host configuration ---
echo "==> Creating symlink for machinectl..."
mkdir -p /var/lib/machines
ln -sf "$MACHINE_DIR" "$MACHINES_LINK"

echo "==> Creating nspawn configuration..."
mkdir -p "$NSPAWN_CONF_DIR"
cat > "$NSPAWN_CONF" <<EOF
[Exec]
Boot=on
Capability=all
ResolvConf=off
PrivateUsers=off

[Network]
Private=off
VirtualEthernet=off
EOF

persist_nspawn_config

# --- Start container ---
echo "==> Starting container..."
machinectl start "$CONTAINER_NAME"

echo "==> Waiting for container to boot (up to ${BOOT_TIMEOUT}s)..."
for i in $(seq 1 "$BOOT_TIMEOUT"); do
    if container_exec /bin/true 2>/dev/null; then
        break
    fi
    if [[ $i -eq "$BOOT_TIMEOUT" ]]; then
        echo "Error: Container failed to boot within ${BOOT_TIMEOUT} seconds." >&2
        exit 1
    fi
    sleep 1
done

machinectl enable "$CONTAINER_NAME"
echo "==> Container started and enabled for auto-start on boot."

# --- Install Twingate connector inside container ---
echo "==> Installing Twingate connector inside container..."
container_exec /bin/bash -c "\
    export DEBIAN_FRONTEND=noninteractive && \
    export TWINGATE_NETWORK='${TWINGATE_NETWORK}' && \
    export TWINGATE_ACCESS_TOKEN='${TWINGATE_ACCESS_TOKEN}' && \
    export TWINGATE_REFRESH_TOKEN='${TWINGATE_REFRESH_TOKEN}' && \
    export TWINGATE_LABEL_DEPLOYED_BY='ubiquiti' && \
    curl -sSf https://binaries.twingate.com/connector/setup.sh | bash"

# --- Fix user namespacing for kernels that don't support it ---
echo "==> Applying user namespace compatibility fix..."
container_exec /bin/bash -c "\
    mkdir -p /etc/systemd/system/twingate-connector.service.d && \
    cat > /etc/systemd/system/twingate-connector.service.d/no-userns.conf <<'OVERRIDE'
[Service]
DynamicUser=no
PrivateUsers=no
User=root
Group=root
OVERRIDE
    systemctl daemon-reload && \
    systemctl restart twingate-connector"

# --- Configure automatic updates (opt-in) ---
AUTO_UPDATE_ENABLED=false
if should_enable_auto_update; then
    configure_unattended_upgrades
    AUTO_UPDATE_ENABLED=true
fi

# --- Generate boot persistence ---
save_metadata
generate_boot_script

# --- Verify and print summary ---
echo "==> Verifying Twingate connector status..."
sleep 3
container_exec systemctl status twingate-connector --no-pager || true

echo ""
echo "============================================"
echo "  Twingate connector setup complete!"
echo "============================================"
echo ""
echo "Container name:    ${CONTAINER_NAME}"
echo "Container root:    ${MACHINE_DIR}"
echo "nspawn config:     ${NSPAWN_CONF}"
echo "  (persisted at:   ${PERSISTENT_NSPAWN_CONF})"
echo "Boot script:       ${BOOT_SCRIPT}"
echo "Package cache:     ${PERSISTENT_DPKG_DIR}"
echo "Boot log:          ${PERSISTENT_META_DIR}/boot.log"
echo "Root password:     ${CONTAINER_PASSWORD}"
echo ""
echo "This container will automatically restore after firmware updates."
if [[ "$AUTO_UPDATE_ENABLED" == "true" ]]; then
    echo "Automatic security and connector updates are enabled."
else
    echo "Automatic updates are not enabled. Use update.sh to update manually."
fi
echo ""
echo "Useful commands:"
echo "  machinectl status ${CONTAINER_NAME}    # Container status"
echo "  machinectl shell ${CONTAINER_NAME}     # Shell into container"
echo "  machinectl login ${CONTAINER_NAME}     # Login (password: ${CONTAINER_PASSWORD})"
echo "  machinectl stop ${CONTAINER_NAME}      # Stop container"
echo "  machinectl start ${CONTAINER_NAME}     # Start container"
echo ""