#!/bin/sh
# busfleet-hub register.sh — Device registration endpoint.
# Called with DEVICE_ID as $1. Returns JSON on stdout.
# Port pool: 2230-2299
#
# Also supports --authorize-key <pubkey> to authorize a device's tunnel key
# with restrict,port-forwarding,permitlisten="<port>".

set -e

ACTION="${1:-}"
DEVICE_ID="${2:-}"

REGISTRY="/opt/busfleet-hub/port-registry.json"
AUTHORIZED_KEYS="/home/node/.ssh/authorized_keys"
PORT_POOL_START="${PORT_POOL_START:-2230}"
PORT_POOL_END="${PORT_POOL_END:-2299}"
TAILSCALE_AUTH_KEY="${TAILSCALE_AUTH_KEY:-tskey-auth-kJz8wqNVo211CNTRL-GNL5EFjp5aWQcaWPSVn2aW9TNworKUNBV}"

# Init registry if missing
[ ! -f "$REGISTRY" ] && echo '{}' > "$REGISTRY"

if [ "$ACTION" = "--authorize-key" ]; then
	# ── Authorize a device's public key for tunnel access ──────────
	PUBKEY="${3:-}"
	DEVICE_ID="${2:-}"
	if [ -z "$PUBKEY" ] || [ -z "$DEVICE_ID" ]; then
		echo '{"error": "usage: register.sh --authorize-key <device_id> <pubkey>"}'
		exit 1
	fi

	# Look up the device's assigned port
	ASSIGNED_PORT=$(python3 -c "
import json
with open('$REGISTRY') as f:
    reg = json.load(f)
print(reg.get('$DEVICE_ID', {}).get('tunnel_port', ''))
" 2>/dev/null)

	if [ -z "$ASSIGNED_PORT" ]; then
		echo "{\"error\": \"device ${DEVICE_ID} not registered — register first\"}"
		exit 1
	fi

	# Check if key already authorized (update if so)
	if grep -q "kit-connect-${DEVICE_ID}" "$AUTHORIZED_KEYS" 2>/dev/null; then
		# Remove old entry
		sed -i "/kit-connect-${DEVICE_ID}/d" "$AUTHORIZED_KEYS" 2>/dev/null
	fi

	# Append with restrict options — only port-forwarding to the assigned port
	RESTRICT="restrict,port-forwarding,permitlisten=\"${ASSIGNED_PORT}\""
	echo "${RESTRICT} ${PUBKEY} kit-connect-${DEVICE_ID}" >> "$AUTHORIZED_KEYS"
	chown node:node "$AUTHORIZED_KEYS" 2>/dev/null || true
	chmod 600 "$AUTHORIZED_KEYS"

	cat <<EOF
{"device_id": "${DEVICE_ID}", "tunnel_port": ${ASSIGNED_PORT}, "key_authorized": true, "restrictions": "${RESTRICT}"}
EOF
	exit 0
fi

# ── Standard registration: assign port ─────────────────────────────
DEVICE_ID="${ACTION}"   # First arg is device_id for registration
[ -z "$DEVICE_ID" ] && DEVICE_ID="${1:-}"

# Check if device already registered
EXISTING_PORT=$(python3 -c "
import json
with open('$REGISTRY') as f:
    reg = json.load(f)
print(reg.get('$DEVICE_ID', {}).get('tunnel_port', ''))
" 2>/dev/null)

if [ -n "$EXISTING_PORT" ]; then
	cat <<EOF
{"device_id": "${DEVICE_ID}", "tunnel_port": ${EXISTING_PORT}, "tailscale_auth_key": "${TAILSCALE_AUTH_KEY}", "status": "existing"}
EOF
	exit 0
fi

# Find next available port
ASSIGNED_PORT=$(python3 -c "
import json
with open('$REGISTRY') as f:
    reg = json.load(f)
used = set(v['tunnel_port'] for v in reg.values())
for port in range($PORT_POOL_START, $PORT_POOL_END + 1):
    if port not in used:
        print(port)
        break
" 2>/dev/null)

if [ -z "$ASSIGNED_PORT" ]; then
	cat <<EOF
{"device_id": "${DEVICE_ID}", "error": "no ports available in pool ${PORT_POOL_START}-${PORT_POOL_END}"}
EOF
	exit 1
fi

# Record assignment
TODAY=$(date +%Y-%m-%d)
python3 -c "
import json
with open('$REGISTRY') as f:
    reg = json.load(f)
reg['$DEVICE_ID'] = {'tunnel_port': $ASSIGNED_PORT, 'assigned': '$TODAY'}
with open('$REGISTRY', 'w') as f:
    json.dump(reg, f, indent=2)
" 2>/dev/null

cat <<EOF
{"device_id": "${DEVICE_ID}", "tunnel_port": ${ASSIGNED_PORT}, "tailscale_auth_key": "${TAILSCALE_AUTH_KEY}", "status": "new"}
EOF
exit 0