#!/bin/bash
# setup-ssh-access.sh - Provision SSH access to all LXC containers
#
# This script:
# 1. Ensures openssh-server is installed in each container
# 2. Creates /root/.ssh directory with correct permissions
# 3. Adds the workstation public key to authorized_keys
# 4. Configures PermitRootLogin with key-only authentication
# 5. Starts and enables sshd
#
# Usage: ./setup-ssh-access.sh [vmid...]
#   Without arguments: provisions all containers
#   With arguments: provisions only specified VMIDs

set -euo pipefail

WORKSTATION_KEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIPtmU1h0wIQiIF0UajcUKV4wQQ4f3dFIAHV8j9pQlNnT kavren@KavDesktop-Cachy"

# Map VMIDs to their host nodes
declare -A CONTAINER_NODES=(
    [101]="pm1"    # twingate
    [102]="pm1"    # zwave-js-ui
    [103]="pm4"    # pihole
    [104]="pm2"    # traefik
    [105]="pm2"    # sonarr
    [106]="pm3"    # mqtt
    [107]="pm3"    # dockge
    [108]="pm2"    # radarr
    [110]="pm4"    # docker-pm4
    [111]="pm4"    # unifi
    [112]="pm3"    # foundryvtt
    [113]="pm2"    # docker-pm2
    [114]="pm2"    # prowlarr
    [115]="pm2"    # jellyseerr
    [116]="pm2"    # authelia
    [117]="pm2"    # whisparr
    [118]="pm2"    # notifiarr
    [119]="pm2"    # bazarr
    [120]="pm2"    # kometa
    [121]="elantris" # jellyfin
    [122]="pm2"    # recyclarr
    [123]="elantris" # ollama
    [124]="elantris" # amp
    [125]="pm4"    # vaultwarden
    [126]="pm4"    # immich
    [127]="pm4"    # gitea
    [128]="pm3"    # frigate
)

setup_ssh() {
    local vmid=$1
    local node=${CONTAINER_NODES[$vmid]:-}

    if [[ -z "$node" ]]; then
        echo "ERROR: Unknown VMID $vmid"
        return 1
    fi

    local hostname
    hostname=$(ssh "$node" "pct config $vmid 2>/dev/null | grep hostname | awk '{print \$2}'" 2>/dev/null || echo "unknown")

    echo "=== Setting up SSH for VMID $vmid ($hostname) on $node ==="

    # Check if container is running
    if ! ssh "$node" "pct status $vmid 2>/dev/null" | grep -q "running"; then
        echo "  WARNING: Container $vmid is not running, skipping..."
        return 1
    fi

    # Detect package manager and install openssh-server
    echo "  Installing openssh-server..."
    ssh "$node" "pct exec $vmid -- bash -c '
        if command -v apt-get &> /dev/null; then
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq && apt-get install -y -qq openssh-server 2>/dev/null
        elif command -v apk &> /dev/null; then
            apk add --quiet openssh openssh-server 2>/dev/null
        elif command -v dnf &> /dev/null; then
            dnf install -y -q openssh-server 2>/dev/null
        elif command -v pacman &> /dev/null; then
            pacman -Sy --noconfirm openssh 2>/dev/null
        else
            echo \"Unknown package manager\"
            exit 1
        fi
    '" 2>/dev/null || echo "  Note: openssh may already be installed"

    # Create .ssh directory and set permissions
    echo "  Configuring SSH keys..."
    ssh "$node" "pct exec $vmid -- bash -c '
        mkdir -p /root/.ssh
        chmod 700 /root/.ssh
        touch /root/.ssh/authorized_keys
        chmod 600 /root/.ssh/authorized_keys
    '"

    # Add the workstation key (idempotent - only adds if not present)
    ssh "$node" "pct exec $vmid -- bash -c \"
        grep -qF '$WORKSTATION_KEY' /root/.ssh/authorized_keys 2>/dev/null || \
        echo '$WORKSTATION_KEY' >> /root/.ssh/authorized_keys
    \""

    # Enable PermitRootLogin with key only (more secure than password)
    echo "  Configuring sshd..."
    ssh "$node" "pct exec $vmid -- bash -c '
        if [[ -f /etc/ssh/sshd_config ]]; then
            sed -i \"s/^#*PermitRootLogin.*/PermitRootLogin prohibit-password/\" /etc/ssh/sshd_config
        fi
    '" 2>/dev/null || true

    # Start and enable sshd
    echo "  Starting sshd..."
    ssh "$node" "pct exec $vmid -- bash -c '
        if command -v systemctl &> /dev/null; then
            systemctl enable ssh 2>/dev/null || systemctl enable sshd 2>/dev/null || true
            systemctl restart ssh 2>/dev/null || systemctl restart sshd 2>/dev/null || true
        elif command -v rc-service &> /dev/null; then
            rc-update add sshd default 2>/dev/null || true
            rc-service sshd restart 2>/dev/null || true
        fi
    '" 2>/dev/null || true

    echo "  SSH setup complete for $vmid ($hostname)"
    echo ""
}

# Main execution
if [[ $# -gt 0 ]]; then
    # Provision specific VMIDs
    CONTAINERS=("$@")
else
    # Provision all containers
    CONTAINERS=(${!CONTAINER_NODES[@]})
fi

echo "Starting SSH provisioning for ${#CONTAINERS[@]} containers..."
echo "Using key: $WORKSTATION_KEY"
echo ""

failed=()
for vmid in "${CONTAINERS[@]}"; do
    if ! setup_ssh "$vmid"; then
        failed+=("$vmid")
    fi
done

echo "=== SSH Provisioning Complete ==="
if [[ ${#failed[@]} -gt 0 ]]; then
    echo "Failed containers: ${failed[*]}"
fi
echo ""
echo "Test with: ssh root@<container-ip>"
echo "Or after DNS setup: ssh root@<service>.kav"