#!/bin/sh
# shellcheck shell=sh
#
# zt-gateway-setup — configure this OpenWrt router as a ZeroTier exit
# gateway for the WIBLAN subnet (10.11.13.0/24).
#
# Subcommands:
#   status          Show current setup state
#   setup-bridge    Create the br-zt bridge interface
#   setup-routing   Configure policy routing (tables 100/101, ip rules)
#   setup-dhcp      Configure DHCP for WIBLAN on br-zt
#   setup-hotplug   Create hotplug script to re-apply routes on ifup
#   setup-persistence  Write rc.local + UCI network routes
#   setup-all       Run all setup-* commands in order
#
# Exit codes:
#   0  success
#   1  usage / argument error
#   2  precondition failed
#   3  runtime failure
#
# Environment overrides (for testing and non-default configs):
#   ZTG_BRIDGE           bridge device            (default: br-zt)
#   ZTG_BRIDGE_PORTS     space-separated ports     (default: ztabc0)
#   ZTG_WIBLAN_CIDR      WIBLAN subnet            (default: 10.11.13.0/24)
#   ZTG_WIBLAN_GW        WIBLAN gateway IP        (default: 10.11.13.1)
#   ZTG_WIBLAN_LEASE_FIRST first DHCP IP          (default: 10.11.13.100)
#   ZTG_WIBLAN_LEASE_LAST  last DHCP IP           (default: 10.11.13.200)
#   ZTG_TABLE_MAIN       main policy table        (default: 100)
#   ZTG_TABLE_DRAIN      drain policy table       (default: 101)
#   ZTG_TABLE_MWAN       mwan3 return table       (default: 1)
#   ZTG_FWMARK           drain fwmark             (default: 0x100)
#   ZTG_DRAIN_PRIORITY   drain rule priority      (default: 99)
#   ZTG_HOTPLUG          hotplug script path
#   ZTG_RCLOCAL          rc.local path
#   ZTG_DHCPCONF         DHCP UCI config file     (default: /etc/config/dhcp)
#   ZTG_NETWORKCONF      network UCI config file  (default: /etc/config/network)
#   ZTG_SKIP_PERSIST     skip UCI persistence (testing)

set -eu

# ----------------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------------
BRIDGE="${ZTG_BRIDGE:-br-zt}"
BRIDGE_PORTS="${ZTG_BRIDGE_PORTS:-ztabc0}"
WIBLAN_CIDR="${ZTG_WIBLAN_CIDR:-10.11.13.0/24}"
WIBLAN_GW="${ZTG_WIBLAN_GW:-10.11.13.1}"
WIBLAN_LEASE_FIRST="${ZTG_WIBLAN_LEASE_FIRST:-10.11.13.100}"
WIBLAN_LEASE_LAST="${ZTG_WIBLAN_LEASE_LAST:-10.11.13.200}"
TABLE_MAIN="${ZTG_TABLE_MAIN:-100}"
TABLE_DRAIN="${ZTG_TABLE_DRAIN:-101}"
TABLE_MWAN="${ZTG_TABLE_MWAN:-1}"
FWMARK="${ZTG_FWMARK:-0x100}"
DRAIN_PRIORITY="${ZTG_DRAIN_PRIORITY:-99}"
HOTPLUG="${ZTG_HOTPLUG:-/etc/hotplug.d/net/99-zerotier-bridge}"
RCLOCAL="${ZTG_RCLOCAL:-/etc/rc.local}"
DHCPCONF="${ZTG_DHCPCONF:-/etc/config/dhcp}"
SKIP_PERSIST="${ZTG_SKIP_PERSIST:-0}"

# Derived: extract prefix bits from CIDR
WIBLAN_BITS="${WIBLAN_CIDR##*/}"
# UCI-safe section name (replace hyphens with underscores)
BRIDGE_UCI=$(printf '%s' "$BRIDGE" | tr '-' '_')

# ----------------------------------------------------------------------------
# Logging
# ----------------------------------------------------------------------------
log() { printf '[zt-gateway-setup] %s\n' "$*" >&2; }
die() { rc=$1; shift; log "ERROR: $*"; exit "$rc"; }
# Ensure a UCI config file exists (touch it if missing)
ensure_uci_config() {
	_conf=/etc/config/"$1"
	if [ ! -f "$_conf" ]; then
		touch "$_conf"
		log "created empty UCI config: ${_conf}"
	fi
}


# CIDR to dotted mask (e.g. 24 -> 255.255.255.0)
_cidr_to_mask() {
	bits=$1
	mask=""
	while [ "$bits" -gt 0 ]; do
		if [ "$bits" -ge 8 ]; then
			oct=255
			bits=$((bits - 8))
		else
			# Build partial octet: bits leading 1s in MSB position
			oct=0
			j=0
			while [ $j -lt "$bits" ]; do
				oct=$(( oct | (1 << (7 - j)) ))
				j=$((j + 1))
			done
			bits=0
		fi
		if [ -n "$mask" ]; then
			mask="${mask}.${oct}"
		else
			mask="${oct}"
		fi
	done
	printf '%s' "$mask"
}

# ----------------------------------------------------------------------------
# Status
# ----------------------------------------------------------------------------
cmd_status() {
	log "checking setup status..."

	# Bridge
	if ip link show "$BRIDGE" >/dev/null 2>&1; then
		printf 'bridge:       %s (up)\n' "$BRIDGE"
		# List ports
		ports=$(ip link show master "$BRIDGE" 2>/dev/null \
			| awk -F': ' '/^[0-9]+:/{gsub(/@.*/, "", $2); print $2}' \
			| tr '\n' ' ')
		if [ -n "$ports" ]; then
			printf '  ports:      %s\n' "$ports"
		fi
		# IP on bridge
		br_addr=$(ip -4 addr show dev "$BRIDGE" 2>/dev/null \
			| awk '/inet /{gsub(/\/.*/, "", $2); print $2; exit}')
		if [ -n "$br_addr" ]; then
			printf '  addr:       %s\n' "$br_addr"
		fi
	else
		printf 'bridge:       %s (missing)\n' "$BRIDGE"
	fi

	# Routing tables
	table_main_gw=$(ip route show table "$TABLE_MAIN" 2>/dev/null \
		| awk '/^[[:space:]]*default/{
			for (i=1; i<=NF; i++) if ($i=="via") { print $(i+1); exit }
		}')
	if [ -n "$table_main_gw" ]; then
		printf 'table %s:     default via %s (configured)\n' "$TABLE_MAIN" "$table_main_gw"
	else
		printf 'table %s:     (empty)\n' "$TABLE_MAIN"
	fi

	table_drain_gw=$(ip route show table "$TABLE_DRAIN" 2>/dev/null \
		| awk '/^[[:space:]]*default/{
			for (i=1; i<=NF; i++) if ($i=="via") { print $(i+1); exit }
		}')
	if [ -n "$table_drain_gw" ]; then
		printf 'table %s:     default via %s (configured)\n' "$TABLE_DRAIN" "$table_drain_gw"
	else
		printf 'table %s:     (empty)\n' "$TABLE_DRAIN"
	fi

	# ip rule for fwmark
	if ip rule show 2>/dev/null | grep -q "fwmark ${FWMARK}"; then
		printf 'ip rule:      fwmark %s -> table %s (configured)\n' "$FWMARK" "$TABLE_DRAIN"
	else
		printf 'ip rule:      fwmark %s (missing)\n' "$FWMARK"
	fi

	# DHCP
	if grep -q "interface '${BRIDGE}'" "$DHCPCONF" 2>/dev/null; then
		printf 'dhcp:         %s configured in %s\n' "$BRIDGE" "$DHCPCONF"
	else
		printf 'dhcp:         %s (not configured)\n' "$BRIDGE"
	fi

	# Hotplug
	if [ -x "$HOTPLUG" ] || [ -f "$HOTPLUG" ]; then
		printf 'hotplug:      %s (present)\n' "$HOTPLUG"
	else
		printf 'hotplug:      %s (missing)\n' "$HOTPLUG"
	fi

	# Persistence in rc.local
	if grep -q "zt-gateway" "$RCLOCAL" 2>/dev/null; then
		printf 'rc.local:     entries present\n'
	else
		printf 'rc.local:     no zt-gateway entries\n'
	fi
}

# ----------------------------------------------------------------------------
# setup-bridge
# ----------------------------------------------------------------------------
cmd_setup_bridge() {
	log "setting up bridge ${BRIDGE}..."

	if ! command -v uci >/dev/null 2>&1; then
		die 3 "uci not found; cannot configure bridge"
	fi
	ensure_uci_config network

	# Create or update bridge device in network.uci
	# Note: UCI section names cannot contain hyphens, so we use BRIDGE_UCI
	if uci -q get "network.${BRIDGE_UCI}" >/dev/null 2>&1; then
		log "bridge device ${BRIDGE} already exists in UCI; updating"
	else
		uci -q set "network.${BRIDGE_UCI}=device"
		uci -q set "network.${BRIDGE_UCI}.type=bridge"
		uci -q set "network.${BRIDGE_UCI}.name=${BRIDGE}"
	fi

	# Set bridge ports (space-separated in UCI list)
	uci -q delete "network.${BRIDGE_UCI}.ports" 2>/dev/null || true
	for port in $BRIDGE_PORTS; do
		uci -q add_list "network.${BRIDGE_UCI}.ports=${port}"
	done

	# Create interface section bridging to br-zt for WIBLAN
	if ! uci -q get "network.zt_wiblan" >/dev/null 2>&1; then
		uci -q set "network.zt_wiblan=interface"
		uci -q set "network.zt_wiblan.proto='static'"
		uci -q set "network.zt_wiblan.device='${BRIDGE}'"
		uci -q set "network.zt_wiblan.ipaddr='${WIBLAN_GW}'"
		uci -q set "network.zt_wiblan.netmask='$(_cidr_to_mask "$WIBLAN_BITS")'"
	fi

	if [ "$SKIP_PERSIST" != "1" ]; then
		uci commit network
		log "bridge UCI config committed"
	fi

	# Bring up the bridge (best-effort; may need netifd restart)
	if command -v ifup >/dev/null 2>&1; then
		ifup "zt_wiblan" 2>/dev/null || \
			log "warning: ifup zt_wiblan failed; may need 'service network restart'"
	fi

	log "bridge ${BRIDGE} setup complete"
}

# ----------------------------------------------------------------------------
# setup-routing
# ----------------------------------------------------------------------------
cmd_setup_routing() {
	log "setting up routing (tables ${TABLE_MAIN}/${TABLE_DRAIN})..."

	# Host route to WIBLAN gateway via bridge
	ip route replace "$WIBLAN_GW" dev "$BRIDGE" 2>/dev/null || \
		log "warning: host route to ${WIBLAN_GW} failed"

	# Table 100 (main policy): default via WIBLAN_GW
	ip route replace default via "$WIBLAN_GW" dev "$BRIDGE" table "$TABLE_MAIN"

	# Table 101 (drain): default via WIBLAN_GW (same default; drain overrides per-flow)
	ip route replace default via "$WIBLAN_GW" dev "$BRIDGE" table "$TABLE_DRAIN"

	# mwan3 return-traffic table: route WIBLAN back through bridge
	ip route replace "$WIBLAN_CIDR" dev "$BRIDGE" table "$TABLE_MWAN"

	# ip rule: fwmark 0x100 -> drain table
	ip rule add fwmark "$FWMARK" table "$TABLE_DRAIN" priority "$DRAIN_PRIORITY" 2>/dev/null || \
		ip rule replace fwmark "$FWMARK" table "$TABLE_DRAIN" priority "$DRAIN_PRIORITY"

	log "routing setup complete"

	# Persist to UCI
	if [ "$SKIP_PERSIST" != "1" ] && command -v uci >/dev/null 2>&1; then
		ensure_uci_config network
		# Host route
		if ! uci -q get "network.zt_wiblan_host" >/dev/null 2>&1; then
			uci -q set "network.zt_wiblan_host=route"
			uci -q set "network.zt_wiblan_host.target='${WIBLAN_GW}'"
			uci -q set "network.zt_wiblan_host.interface='${BRIDGE}'"
		fi

		# Main policy table default route
		if ! uci -q get "network.zt_wiblan_default" >/dev/null 2>&1; then
			uci -q set "network.zt_wiblan_default=route"
			uci -q set "network.zt_wiblan_default.target='0.0.0.0'"
			uci -q set "network.zt_wiblan_default.netmask='0.0.0.0'"
			uci -q set "network.zt_wiblan_default.gateway='${WIBLAN_GW}'"
			uci -q set "network.zt_wiblan_default.interface='${BRIDGE}'"
			uci -q set "network.zt_wiblan_default.table='${TABLE_MAIN}'"
		fi

		uci commit network
		log "routing UCI config committed"
	fi
}

# ----------------------------------------------------------------------------
# setup-dhcp
# ----------------------------------------------------------------------------
cmd_setup_dhcp() {
	log "setting up DHCP for ${BRIDGE} (${WIBLAN_CIDR})..."

	if ! command -v uci >/dev/null 2>&1; then
		die 3 "uci not found; cannot configure DHCP"
	fi
	ensure_uci_config dhcp

	# Create DHCP subnet entry for br-zt
	# Note: UCI section names cannot contain hyphens, so we use BRIDGE_UCI
	if uci -q get "dhcp.${BRIDGE_UCI}" >/dev/null 2>&1; then
		log "DHCP entry for ${BRIDGE} already exists; updating"
	else
		uci -q set "dhcp.${BRIDGE_UCI}=dhcp"
	fi

	uci -q set "dhcp.${BRIDGE_UCI}.interface=${BRIDGE}"
	uci -q set "dhcp.${BRIDGE_UCI}.start=${WIBLAN_LEASE_FIRST##*.}"
	uci -q set "dhcp.${BRIDGE_UCI}.limit=$(( ${WIBLAN_LEASE_LAST##*.} - ${WIBLAN_LEASE_FIRST##*.} + 1 ))"
	uci -q set "dhcp.${BRIDGE_UCI}.leasetime=12h"

	# Ignore WIBLAN subnet in upstream DHCP (prevent handing out
	# conflicting leases on the LAN side)
	lan_iface=$(uci -q get dhcp.lan.interface 2>/dev/null || echo "lan")
	if [ -n "$lan_iface" ]; then
		# Add WIBLAN to lan's ignore list if not already there
		if ! uci -q get "dhcp.lan.ignore" 2>/dev/null | grep -q "$WIBLAN_CIDR"; then
			uci -q add_list "dhcp.lan.dhcp_option='6,${WIBLAN_GW}'" 2>/dev/null || true
		fi
	fi

	if [ "$SKIP_PERSIST" != "1" ]; then
		uci commit dhcp
		log "DHCP UCI config committed"
	fi

	# Restart dnsmasq to pick up changes
	if command -v service >/dev/null 2>&1; then
		service dnsmasq restart 2>/dev/null || \
			log "warning: dnsmasq restart failed; do it manually"
	fi

	log "DHCP setup complete"
}

# ----------------------------------------------------------------------------
# setup-hotplug
# ----------------------------------------------------------------------------
cmd_setup_hotplug() {
	log "creating hotplug script ${HOTPLUG}..."

	mkdir -p "$(dirname "$HOTPLUG")"

	cat >"$HOTPLUG" <<'HOTPLUG_SCRIPT'
#!/bin/sh
# shellcheck shell=sh
# ZeroTier gateway hotplug — re-apply routing when the bridge interface
# comes up (e.g. after boot, after ZeroTier restart).
#
# Environment: INTERFACE, ACTION (set by netifd hotplug)
# Config: /etc/config/zt-gateway (read at runtime for active gateway IP)

ZTG_BRIDGE="${ZTG_BRIDGE:-br-zt}"
ZTG_TABLE_MAIN="${ZTG_TABLE_MAIN:-100}"
ZTG_TABLE_DRAIN="${ZTG_TABLE_DRAIN:-101}"
ZTG_TABLE_MWAN="${ZTG_TABLE_MWAN:-1}"
ZTG_FWMARK="${ZTG_FWMARK:-0x100}"
ZTG_DRAIN_PRIORITY="${ZTG_DRAIN_PRIORITY:-99}"
ZTG_WIBLAN_CIDR="${ZTG_WIBLAN_CIDR:-10.11.13.0/24}"

# Only act on our bridge interface
[ "$INTERFACE" = "$ZTG_BRIDGE" ] || exit 0

case "$ACTION" in
	ifup)
		# Read active gateway from UCI
		active_ip=$(uci -q get zt-gateway.global.active_ip 2>/dev/null || \
			uci -q get zt-gateway.global.active_gateway 2>/dev/null)

		# If active_ip is a region name, resolve to IP via gateway section
		if [ -n "$active_ip" ] && ! echo "$active_ip" | grep -q '^[0-9]'; then
			active_ip=$(uci -q get "zt-gateway.@gateway[0].ip" 2>/dev/null || \
				for sec in $(uci -q show zt-gateway 2>/dev/null | \
					awk -F'=' '/\.region=/{print $1}' | \
					sed 's/\.region//'); do
					r=$(uci -q get "${sec}.region" 2>/dev/null)
					if [ "$r" = "$active_ip" ]; then
						uci -q get "${sec}.ip" 2>/dev/null
						break
					fi
				done)
		fi

		if [ -z "$active_ip" ]; then
			# No active gateway configured; try reading from table
			active_ip=$(ip route show table "$ZTG_TABLE_MAIN" 2>/dev/null \
				| awk '/^[[:space:]]*default/{
					for (i=1; i<=NF; i++) if ($i=="via") { print $(i+1); exit }
				}')
		fi

		if [ -z "$active_ip" ]; then
			logger -t zt-gw-hotplug "No active gateway IP; skipping route setup"
			exit 0
		fi

		logger -t zt-gw-hotplug "ifup ${ZTG_BRIDGE}: applying routes via ${active_ip}"

		# Host route
		ip route replace "$active_ip" dev "$ZTG_BRIDGE"

		# Policy routes
		ip route replace default via "$active_ip" dev "$ZTG_BRIDGE" table "$ZTG_TABLE_MAIN"
		ip route replace default via "$active_ip" dev "$ZTG_BRIDGE" table "$ZTG_TABLE_DRAIN"

		# mwan3 return
		ip route replace "$ZTG_WIBLAN_CIDR" dev "$ZTG_BRIDGE" table "$ZTG_TABLE_MWAN"

		# ip rule for drain fwmark
		ip rule add fwmark "$ZTG_FWMARK" table "$ZTG_TABLE_DRAIN" \
			priority "$ZTG_DRAIN_PRIORITY" 2>/dev/null || \
			ip rule replace fwmark "$ZTG_FWMARK" table "$ZTG_TABLE_DRAIN" \
				priority "$ZTG_DRAIN_PRIORITY"
		;;
	ifdown)
		logger -t zt-gw-hotplug "ifdown ${ZTG_BRIDGE}: cleaning up"
		ip rule del fwmark "$ZTG_FWMARK" table "$ZTG_TABLE_DRAIN" \
			priority "$ZTG_DRAIN_PRIORITY" 2>/dev/null || true
		;;
esac
HOTPLUG_SCRIPT

	chmod +x "$HOTPLUG"
	log "hotplug script created at ${HOTPLUG}"
}

# ----------------------------------------------------------------------------
# setup-persistence
# ----------------------------------------------------------------------------
cmd_setup_persistence() {
	log "setting up persistence..."

	# Ensure rc.local has the gateway restoration logic
	if ! grep -q "zt-gateway" "$RCLOCAL" 2>/dev/null; then
		log "adding zt-gateway entry to ${RCLOCAL}"

		# Read current rc.local content
		rc_content=""
		if [ -f "$RCLOCAL" ]; then
			rc_content=$(cat "$RCLOCAL")
		fi

		# Remove trailing 'exit 0' if present, add our block, re-add exit 0
		rc_stripped=$(printf '%s\n' "$rc_content" | sed '/^exit 0$/d')

		cat >"$RCLOCAL" <<RCEOF
${rc_stripped}
# --- zt-gateway: restore routes on boot ---
# Applied via /etc/hotplug.d/net/99-zerotier-bridge on ifup
# This block ensures the bridge comes up at boot
[ -x /etc/init.d/network ] && /etc/init.d/network reload

exit 0
RCEOF
		log "rc.local updated"
	else
		log "rc.local already contains zt-gateway entries; skipping"
	fi

	# Write active gateway IP to UCI for the hotplug script
	if [ "$SKIP_PERSIST" != "1" ] && command -v uci >/dev/null 2>&1; then
		ensure_uci_config zt-gateway
		# Store the WIBLAN config for the hotplug to read
		if ! uci -q get "zt-gateway.global.wiblan_subnet" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.wiblan_subnet='${WIBLAN_CIDR}'"
		fi
		if ! uci -q get "zt-gateway.global.bridge_device" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.bridge_device='${BRIDGE}'"
		fi
		if ! uci -q get "zt-gateway.global.bridge_ports" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.bridge_ports='${BRIDGE_PORTS}'"
		fi
		if ! uci -q get "zt-gateway.global.wiblan_gw" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.wiblan_gw='${WIBLAN_GW}'"
		fi
		if ! uci -q get "zt-gateway.global.table_main" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.table_main='${TABLE_MAIN}'"
		fi
		if ! uci -q get "zt-gateway.global.table_drain" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.table_drain='${TABLE_DRAIN}'"
		fi
		if ! uci -q get "zt-gateway.global.dhcp_lease_first" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.dhcp_lease_first='${WIBLAN_LEASE_FIRST}'"
		fi
		if ! uci -q get "zt-gateway.global.dhcp_lease_last" >/dev/null 2>&1; then
			uci -q set "zt-gateway.global.dhcp_lease_last='${WIBLAN_LEASE_LAST}'"
		fi
		uci commit zt-gateway
		log "UCI global config updated with setup parameters"
	fi

	log "persistence setup complete"
}

# ----------------------------------------------------------------------------
# setup-all
# ----------------------------------------------------------------------------
cmd_setup_all() {
	log "running full gateway setup..."

	cmd_setup_bridge
	cmd_setup_routing
	cmd_setup_dhcp
	cmd_setup_hotplug
	cmd_setup_persistence

	log "========================================="
	log "Full gateway setup complete!"
	log "  Bridge:     ${BRIDGE}"
	log "  Ports:      ${BRIDGE_PORTS}"
	log "  Subnet:     ${WIBLAN_CIDR}"
	log "  Gateway IP: ${WIBLAN_GW}"
	log "  Tables:     ${TABLE_MAIN} (main), ${TABLE_DRAIN} (drain)"
	log "  DHCP:       ${WIBLAN_LEASE_FIRST} - ${WIBLAN_LEASE_LAST}"
	log "========================================="
}

# ----------------------------------------------------------------------------
# Usage
# ----------------------------------------------------------------------------
usage() {
	cat >&2 <<'USAGE'
Usage: zt-gateway-setup <command>

Commands:
  status          Show current setup state
  setup-bridge    Create the br-zt bridge interface
  setup-routing   Configure policy routing (tables 100/101, ip rules)
  setup-dhcp      Configure DHCP for WIBLAN on br-zt
  setup-hotplug   Create hotplug script to re-apply routes on ifup
  setup-persistence  Write rc.local + UCI network routes
  setup-all       Run all setup-* commands in order
USAGE
}

# ----------------------------------------------------------------------------
# Entry point
# ----------------------------------------------------------------------------
main() {
	if [ $# -lt 1 ]; then
		usage
		die 1 "missing command"
	fi

	cmd=$1
	shift

	case "$cmd" in
		status)         cmd_status "$@" ;;
		setup-bridge)   cmd_setup_bridge "$@" ;;
		setup-routing)  cmd_setup_routing "$@" ;;
		setup-dhcp)     cmd_setup_dhcp "$@" ;;
		setup-hotplug)  cmd_setup_hotplug "$@" ;;
		setup-persistence) cmd_setup_persistence "$@" ;;
		setup-all)      cmd_setup_all "$@" ;;
		-h|--help)      usage; exit 0 ;;
		*)              usage; die 1 "unknown command: $cmd" ;;
	esac
}

main "$@"