#!/bin/sh
# omr-weight-sync: Sync UCI multipath_weight → BPF endpoint_weights map
# and uci network.<iface>.weight (ip route semantics).
# When enable_decision_weights=1 (default), polls GET /metrics/decision on
# each configured VPS first and keeps model-assigned weights in /var/state
# (openmptcprouter.<iface>.decision_weight), which then go to the BPF map
# instead of the UCI ones.
# Polls every MPTCP_WEIGHT_INTERVAL seconds (default: 5).

. "${OMR_LIB_DIR:-/usr/share/omr/lib}/omr-vps-curl.sh"
# omr_metrics_curl: the same, and the custom metrics server pinned too
. "${OMR_METRICS_LIB_DIR:-/usr/share/omr/lib}/omr-metrics-curl.sh"
. "${OMR_LIB_DIR:-/usr/share/omr/lib}/omr-state.sh"
POLL_INTERVAL="${MPTCP_WEIGHT_INTERVAL:-5}"
WEIGHT_SCRIPT="/usr/sbin/mptcp-scheduler-weight.sh"
MAP_PATH="/sys/fs/bpf/endpoint_weights"
# The IPv6 endpoint-id map "set" fills too (absent with an older scheduler)
MAP6_PATH="/sys/fs/bpf/weight_local_id"
# Shared with omr-metrics-send (and its login backoff): the .no_decision
# marks and the weights last applied per device live here
SYNC_STATE_DIR="$(_omr_metrics_login_dir)"
APPLIED_DIR="${SYNC_STATE_DIR}/weight-applied"

# _login logged through it, but it was never defined: "_log: not found" on
# stderr and nothing in syslog
_log() {
	logger -t "omr-weight-sync" "$@"
}

# bpf_weight_rr reads the same endpoint_weights map (040-multipath-weight
# feeds it too): without it here, a switch from bpf_weight left the last
# model weight in /var/state, and 040 kept putting it in the map
_is_weight_scheduler() {
	local s
	s=$(uci -q get network.globals.mptcp_scheduler 2>/dev/null)
	case "$s" in
		bpf_weight|bpf_burstweight|bpf_weight_rr|bpf_overflow|mptcp_bpf_weight.o|mptcp_bpf_burstweight.o|mptcp_bpf_weight_rr.o|mptcp_bpf_overflow.o) return 0 ;;
	esac
	return 1
}

# How long the model's weights are kept while no VPS answers
DECISION_TTL="${MPTCP_DECISION_TTL:-300}"
_decision_ok=""
# Whether /var/state may hold model weights: unknown at start
_decision_held=1

_uptime() {
	local _up _rest
	read -r _up _rest < /proc/uptime
	echo "${_up%%.*}"
}

# _weight_bound <weight>: 1 to 256, as "ip route" takes it and LuCI asks;
# nothing for a value that is no number. The schedulers read 0 as no weight
# (neutral 100) and add the weight to a signed counter: a weight set by hand
# outside that range did not do what it said.
_weight_bound() {
	local _w="$1"
	case "$_w" in
		''|*[!0-9]*) return 0 ;;
	esac
	_w="${_w#"${_w%%[!0]*}"}"
	if [ -z "$_w" ]; then
		_w=1
	elif [ "${#_w}" -gt 3 ] || [ "$_w" -gt 256 ]; then
		_w=256
	fi
	echo "$_w"
}

# _clear_decision_weights [<interface> ...]: drop the model's weight of these
# interfaces, of all of them when none is given. 040-multipath-weight,
# 006-mqvpn-weight and the VPS pins read it as long as weight sync and
# decision weights are on: a weight the VPS no longer sends, or one left from
# another scheduler, stayed in use until a reboot.
_clear_decision_weights() {
	local _iface
	# shellcheck disable=SC2046 # one uci section name per word
	[ $# -eq 0 ] && set -- $(uci -q -P "${OMR_STATE_DIR:-/var/state}" show openmptcprouter 2>/dev/null | \
		sed -n 's/^openmptcprouter\.\([^.=][^.=]*\)\.decision_weight=.*/\1/p')
	for _iface in "$@"; do
		omr_state_del openmptcprouter "$_iface" decision_weight
	done
	return 0
}

# _is_ip6, _server_creds, _login (with the login backoff omr-metrics-send and
# the metrics rpcd share), _token_store and _token_drop come from
# omr-metrics-curl.sh. This loop runs every 5 s: its own login, with no
# backoff, sent a wrong key to the API every 5 s.
SRV_TOKEN=""
SRV_TOKEN_KEY=""
SRV_USERNAME=""
SRV_PASSWORD=""

# _no_decision <servername>: the server answered 404 (no omr_metrics module)
# or 501 (no decision engine, PyTorch missing) to /metrics/decision less than
# an hour ago. Without this mark the loop asked it again every 5 s, for good.
# Its own mark, not omr-metrics-send's .no_metrics: a VPS without the decision
# engine still stores the metrics. An hour, like that one: the module can be
# installed on the VPS later.
_no_decision() {
	local f="${SYNC_STATE_DIR}/${1}.no_decision" since
	[ -f "$f" ] || return 1
	since=$(date -r "$f" +%s 2>/dev/null)
	case "$since" in ''|*[!0-9]*) since=0 ;; esac
	[ $(($(date +%s) - since)) -lt 3600 ] && return 0
	rm -f "$f"
	return 1
}

# GET /metrics/decision from one VPS server.
# Args: servername server serverport predict(true|false) horizon(seconds)
# Prints the raw JSON response on stdout; returns 1 on failure.
_get_decision() {
	local servername="$1" server="$2" serverport="$3" predict="$4" horizon="$5"
	local url http_code tmpfile

	_no_decision "$servername" && return 1
	# The omr-metrics credentials are the custom server's only: the VPS
	# rejected them once the custom server was turned off
	_server_creds "$servername"
	[ -z "$SRV_TOKEN" ] && { _login "$servername" "$server" "$serverport" || return 1; }

	if _is_ip6 "$server"; then
		url="https://[${server}]:${serverport}/metrics/decision?predict=${predict}&horizon=${horizon}"
	else
		url="https://${server}:${serverport}/metrics/decision?predict=${predict}&horizon=${horizon}"
	fi

	tmpfile=$(mktemp /tmp/omr-decision.XXXXXX) || return 1

	http_code=$(omr_metrics_curl --max-time 10 -s -k \
		-o "$tmpfile" -w "%{http_code}" \
		-H "accept: application/json" \
		-H "Authorization: Bearer ${SRV_TOKEN}" \
		"$url" 2>/dev/null)

	if [ "$http_code" = "401" ]; then
		# Refused: not sent again every 5 s while the login backs off
		_token_drop "$SRV_TOKEN_KEY"
		SRV_TOKEN=""
		_login "$servername" "$server" "$serverport" && http_code=$(omr_metrics_curl --max-time 10 -s -k \
			-o "$tmpfile" -w "%{http_code}" \
			-H "accept: application/json" \
			-H "Authorization: Bearer ${SRV_TOKEN}" \
			"$url" 2>/dev/null)
	fi

	if [ "$http_code" = "200" ]; then
		cat "$tmpfile"
		# The stored token works: whatever armed a backoff is over
		_login_ok "$(_state_key "$servername" "$server")"
	elif [ "$http_code" = "404" ] || [ "$http_code" = "501" ]; then
		_log "no decision engine on ${server} (HTTP ${http_code}), not asked again for an hour"
		mkdir -p "$SYNC_STATE_DIR" 2>/dev/null
		touch "${SYNC_STATE_DIR}/${servername}.no_decision" 2>/dev/null
	fi
	rm -f "$tmpfile"
	[ "$http_code" = "200" ]
}

# _differs <current> <new>: the new weight is far enough from the current one
# to be applied. The model's weights move by a unit or two from one answer to
# the next.
_differs() {
	local cur="$1" new="$2" d
	case "$cur" in
		''|*[!0-9]*) return 0 ;;
	esac
	d=$((new - cur))
	[ "$d" -lt 0 ] && d=$((-d))
	[ "$d" -ge 2 ] && [ "$((d * 20))" -ge "$cur" ]
}

# _apply_decision <response> <interfaces>: keep the weights of a
# /metrics/decision answer, {"weights": {"<interface>": 1-255, ...}, ...},
# as openmptcprouter.<interface>.decision_weight in /var/state, for
# _apply_weights. They were written to network.<interface>.multipath_weight
# and .weight and committed: a flash write on every change, and the weights
# the user set were lost for good.
_apply_decision() {
	local response="$1" iface_list="$2" _iface _weight _cur
	for _iface in $iface_list; do
		_cur=$(omr_state_get "openmptcprouter.${_iface}.decision_weight")
		_weight=$(jsonfilter -s "$response" -e "@.weights['${_iface}']" 2>/dev/null)
		case "$_weight" in
			''|*[!0-9]*) _weight=0 ;;
		esac
		# No weight for this interface in the answer: the user's applies
		if [ "$_weight" -lt 1 ]; then
			[ -n "$_cur" ] && _clear_decision_weights "$_iface"
			continue
		fi
		[ "$_weight" -gt 255 ] && _weight=255

		_differs "$_cur" "$_weight" && \
			omr_state_set openmptcprouter "$_iface" decision_weight "$_weight"
	done
	return 0
}

# Fetch model-assigned weights from the first responding VPS and keep them
# in /var/state (_apply_decision).  Returns 0 when at least one weight was
# received; falls through to UCI-only weights otherwise.
_fetch_decision_weights() {
	_is_weight_scheduler || return 1

	local custom_server custom_serverport server_names use_custom_server
	use_custom_server=$(uci -q get "omr-metrics.settings.use_custom_server" 2>/dev/null)
	if [ "${use_custom_server:-0}" = "1" ]; then
		custom_server=$(uci -q get "omr-metrics.settings.server" 2>/dev/null)
		custom_serverport=$(uci -q get "omr-metrics.settings.serverport" 2>/dev/null)
		[ -z "$custom_serverport" ] && custom_serverport="65500"
	fi

	server_names=$(uci -q show openmptcprouter 2>/dev/null | \
		sed -n 's/^openmptcprouter\.\([^.=][^.=]*\)=server$/\1/p')
	[ -z "$custom_server" ] && [ -z "$server_names" ] && return 1

	local iface_list
	iface_list=$(uci show network 2>/dev/null \
		| grep -E "multipath='(on|master|backup)'" \
		| sed "s/network\\.\\(.*\\)\\.multipath.*/\\1/")
	[ -z "$iface_list" ] && return 1

	local _predict _horizon _v
	_v=$(uci -q get omr-metrics.settings.decision_predict 2>/dev/null)
	[ "${_v:-0}" = "1" ] && _predict="true" || _predict="false"
	_horizon=$(uci -q get omr-metrics.settings.decision_horizon 2>/dev/null)
	[ -z "$_horizon" ] && _horizon="300"

	# If a custom omr-metrics server is configured, use it exclusively.
	if [ -n "$custom_server" ]; then
		local response
		response=$(_get_decision "omr_metrics_custom" "$custom_server" "$custom_serverport" "$_predict" "$_horizon") || return 1
		_apply_decision "$response" "$iface_list"
		return 0
	fi

	for servername in $server_names; do
		local disabled serverport server_ips
		disabled=$(uci -q get "openmptcprouter.${servername}.disabled" 2>/dev/null)
		[ "$disabled" = "1" ] && continue

		serverport=$(uci -q get "openmptcprouter.${servername}.port" 2>/dev/null)
		[ -z "$serverport" ] && serverport="65500"

		# uci get prints a list space-separated. uci show quotes each value
		# (ip='a' 'b'), and the sed that cut the outer quotes kept the inner
		# ones: with several addresses, every one of them was mangled
		server_ips=$(uci -q get "openmptcprouter.${servername}.ip" 2>/dev/null)

		for server in $server_ips; do
			[ -z "$server" ] && continue

			local response
			response=$(_get_decision "$servername" "$server" "$serverport" "$_predict" "$_horizon") || continue

			_apply_decision "$response" "$iface_list"
			return 0
		done
	done

	return 1
}

# _dev_endpoints <device>: _dev_eps = the MPTCP endpoints of <device> on one
# line, from the list _apply_weights read once for the whole cycle (no
# subshell: this runs for every WAN every 5 s)
_dev_endpoints() {
	local _line
	_dev_eps=""
	while IFS= read -r _line; do
		case "$_line " in *" dev $1 "*) _dev_eps="$_dev_eps$_line;" ;; esac
	done <<-EOF
	$_eps
	EOF
}

# "set" runs bpftool several times per endpoint, and the loop ran it on
# every WAN every 5 s with the same weight. It is called when the weight, the
# device's endpoints ("set" keys the map by endpoint address and id: a new
# DHCP lease needs it again) or the maps themselves (a scheduler reload pins
# new ones) changed since it last succeeded. 040-multipath-weight still
# compares the map itself on every tracker cycle.
_apply_weights() {
	_is_weight_scheduler && [ -e "$MAP_PATH" ] && [ -x "$WEIGHT_SCRIPT" ] || {
		# switched back later, the maps start empty
		[ -d "$APPLIED_DIR" ] && rm -rf "$APPLIED_DIR"
		return 0
	}

	local _tmpfile _changed _iface _dev _weight _mweight _uweight _cur _decision=0
	local _eps _maps _key _applied _afile _dev_eps
	_tmpfile=$(mktemp /tmp/omr-weight.XXXXXX 2>/dev/null) || return 1
	_changed=0
	_decision_enabled && _decision=1
	_eps=$(ip mptcp endpoint show 2>/dev/null)
	_maps=$(ls -di "$MAP_PATH" "$MAP6_PATH" 2>/dev/null | awk '{ printf "%s ", $1 }')
	mkdir -p "$APPLIED_DIR" 2>/dev/null

	uci show network 2>/dev/null \
		| grep -E "multipath='(on|master|backup)'" \
		| sed "s/network\\.\\(.*\\)\\.multipath.*/\\1/" \
		> "$_tmpfile"

	while IFS= read -r _iface; do
		[ -z "$_iface" ] && continue

		# The MPTCP endpoint sits on the l3 device: for PPPoE (or any
		# protocol with its own netdev) the uci device is the underlying
		# port, which has no endpoint, so that WAN never got its weight
		_dev=$(ifstatus "$_iface" 2>/dev/null | jsonfilter -e '@.l3_device' 2>/dev/null)
		[ -z "$_dev" ] && _dev=$(uci -q get "network.${_iface}.device" 2>/dev/null)
		[ -z "$_dev" ] && _dev=$(uci -q get "network.${_iface}.ifname" 2>/dev/null)
		[ -z "$_dev" ] && continue

		# The user's weight: multipath_weight, unset is 100. Not .weight
		# when it is unset: that is the copy made here of the weight before,
		# and a Weight field emptied on the MPTCP page brought it back
		_mweight=$(_weight_bound "$(uci -q get "network.${_iface}.multipath_weight" 2>/dev/null)")
		_uweight="${_mweight:-100}"

		# The model's weight, when there is one, goes to the map instead.
		# It is runtime state (see _apply_decision), never the config.
		_weight=""
		[ "$_decision" = "1" ] && _weight=$(omr_state_get "openmptcprouter.${_iface}.decision_weight")
		case "$_weight" in
			''|*[!0-9]*) _weight="" ;;
		esac

		# Apply to BPF map, when anything it depends on changed
		_dev_endpoints "$_dev"
		_key="${_weight:-$_uweight}|${_maps}|${_dev_eps}"
		# a netdev name has no "/"
		_afile="${APPLIED_DIR}/${_dev}"
		_applied=""
		[ -f "$_afile" ] && read -r _applied < "$_afile"
		if [ "$_applied" != "$_key" ]; then
			if "$WEIGHT_SCRIPT" set "$_dev" "${_weight:-$_uweight}" >/dev/null 2>&1; then
				printf '%s\n' "$_key" > "$_afile"
			else
				rm -f "$_afile"
			fi
		fi

		# Not while a model weight is in the map: 040-multipath-weight
		# leaves .weight alone then too, and the two committed network to
		# flash in turn when they did not
		[ -n "$_weight" ] && continue

		# Sync network.<iface>.weight (ip route semantics) with the
		# MPTCP weight the user set, as 040 does. Only a set one: .weight
		# is also the Weight of the Interfaces page, which the balancing
		# route reads (unset: 100 on the master, 1 on the others), and an
		# unset multipath_weight reset it to 100 every 5 s.
		# multipath_weight stays unset when it is: LuCI shows 100 for it
		# and leaves it unset on save
		[ -z "$_mweight" ] && continue
		_cur=$(uci -q get "network.${_iface}.weight" 2>/dev/null)
		if [ "$_cur" != "$_mweight" ]; then
			uci -q set "network.${_iface}.weight=$_mweight"
			_changed=1
		fi
	done < "$_tmpfile"

	rm -f "$_tmpfile"
	[ "$_changed" -eq 1 ] && uci -q commit network 2>/dev/null || true
}

_decision_enabled() {
	local v
	v=$(uci -q get omr-metrics.settings.enable_decision_weights 2>/dev/null)
	[ "${v:-1}" != "0" ]
}

# _sync_decision: fetch the model's weights while they are wanted, and drop
# them when they are not (decision weights off, no weight scheduler) or when
# no VPS has answered for DECISION_TTL seconds
_sync_decision() {
	local _now
	if _decision_enabled && _is_weight_scheduler; then
		_now=$(_uptime)
		if _fetch_decision_weights; then
			_decision_ok="$_now"
			_decision_held=1
			return 0
		fi
		[ -z "$_decision_ok" ] && _decision_ok="$_now"
		[ "$((_now - _decision_ok))" -lt "$DECISION_TTL" ] && return 0
	fi
	_decision_ok=""
	[ "$_decision_held" = "1" ] || return 0
	_clear_decision_weights
	_decision_held=0
}

# Stopped (service stop, weight sync turned off): the model's weights go
# with it. 040-multipath-weight, 006-mqvpn-weight and the VPS pins read them
# as long as the settings ask for them, and kept the last ones until a reboot.
_on_term() {
	trap - TERM INT
	[ -n "$_sleep_pid" ] && kill "$_sleep_pid" 2>/dev/null
	_clear_decision_weights
	exit 0
}
_sleep_pid=""
trap _on_term TERM INT
# A restart is the way to retry a server marked without a decision engine,
# and to put every weight in the map again
rm -f "${SYNC_STATE_DIR}"/*.no_decision 2>/dev/null
rm -rf "$APPLIED_DIR" 2>/dev/null

while true; do
	_sync_decision
	_apply_weights
	# In the background: the shell runs the trap only once a foreground
	# sleep is over, and procd may have killed it by then
	sleep "$POLL_INTERVAL" &
	_sleep_pid=$!
	wait "$_sleep_pid"
	_sleep_pid=""
done
