#!/bin/sh

[ -n "$1" ] || exit

. /lib/functions.sh

# retrieve args
OMR_QUOTA_INTERFACE="$1"
shift

_PERSIST_DIR="${OMR_QUOTA_STATE_DIR:-/etc/omr-quota/state}"
_TSTATE_DIR="${OMR_QUOTA_THROTTLE_STATE_DIR:-/tmp/omr-quota}"

_get_real_interface() {
	local iface="$1"
	local real
	local _cache="${_TSTATE_DIR}/${iface}.realdev"
	case "$iface" in
		@*)
			real="$(ifstatus "$iface" | jsonfilter -q -e '@["device"]')"
			;;
		*)
			real="$(ifstatus "$iface" | jsonfilter -q -e '@.l3_device')"
			;;
	esac
	# ifstatus reports no l3_device while the interface is administratively
	# down (e.g. cut by this same script for quota enforcement) -- fall back
	# to the last known device name so vnstat usage lookups keep working.
	# Without this, a cut interface reads back as 0 bytes used next cycle,
	# looks like the quota is no longer exceeded, and gets brought back up
	# only to be cut again next cycle: an infinite up/down flap.
	if [ -n "$real" ]; then
		mkdir -p "$_TSTATE_DIR"
		echo "$real" > "$_cache"
	elif [ -f "$_cache" ]; then
		real="$(cat "$_cache")"
	fi
	printf '%s' "$real"
}

_vnstat_usage() {
	local dev="$1"
	local json rx_bytes tx_bytes
	[ -z "$dev" ] && { printf '0 0'; return; }

	if [ -n "$OMR_QUOTA_BEGINDATE" ]; then
		json="$(vnstat -i "$dev" -b "$OMR_QUOTA_BEGINDATE" --json)"
		rx_bytes="$(printf '%s' "$json" | jsonfilter -q -e '@.interfaces[0].traffic.total.rx' | tr -d "\n")"
		tx_bytes="$(printf '%s' "$json" | jsonfilter -q -e '@.interfaces[0].traffic.total.tx' | tr -d "\n")"
	else
		json="$(vnstat -i "$dev" --json)"
		rx_bytes="$(printf '%s' "$json" | jsonfilter -q -e '@.interfaces[0].traffic.month[-1].rx' | tr -d "\n")"
		tx_bytes="$(printf '%s' "$json" | jsonfilter -q -e '@.interfaces[0].traffic.month[-1].tx' | tr -d "\n")"
	fi

	printf '%s %s' "$(( ${rx_bytes:-0} / 1024 ))" "$(( ${tx_bytes:-0} / 1024 ))"
}

_apply_throttle() {
	local dev="$1"
	local dl_kbit=$(( ${OMR_QUOTA_THROTTLE_DL:-1} * 1000 ))
	local ul_kbit=$(( ${OMR_QUOTA_THROTTLE_UL:-1} * 1000 ))
	local ifb_dev="ifb-$(echo "$dev" | tr '/' '-')"
	modprobe ifb 2>/dev/null || true
	if ! ip link show "$ifb_dev" > /dev/null 2>&1; then
		ip link add name "$ifb_dev" type ifb
	fi
	ip link set "$ifb_dev" up
	tc qdisc replace dev "$dev" root tbf rate "${ul_kbit}kbit" burst 32768 latency 400ms
	tc qdisc add dev "$dev" handle ffff: ingress 2>/dev/null || true
	tc filter del dev "$dev" parent ffff: 2>/dev/null || true
	tc filter add dev "$dev" parent ffff: protocol all u32 match u32 0 0 \
		action mirred egress redirect dev "$ifb_dev"
	tc qdisc replace dev "$ifb_dev" root tbf rate "${dl_kbit}kbit" burst 32768 latency 400ms
}

_remove_throttle() {
	local dev="$1"
	local ifb_dev="ifb-$(echo "$dev" | tr '/' '-')"
	tc qdisc del dev "$dev" root 2>/dev/null || true
	tc qdisc del dev "$dev" ingress 2>/dev/null || true
	if ip link show "$ifb_dev" > /dev/null 2>&1; then
		tc qdisc del dev "$ifb_dev" root 2>/dev/null || true
		ip link del "$ifb_dev" 2>/dev/null || true
	fi
}

_apply_downstream_limit() {
	local rate_kbit="$1"
	local iface dev
	[ -z "$rate_kbit" ] && return
	mkdir -p "$_TSTATE_DIR"
	touch "${_TSTATE_DIR}/${OMR_QUOTA_INTERFACE}.downstream"
	for iface in ${OMR_QUOTA_DOWN_INTERFACES:-${OMR_QUOTA_INTERFACES:-$OMR_QUOTA_INTERFACE}}; do
		dev="$(_get_real_interface "$iface")"
		[ -n "$dev" ] && tc qdisc replace dev "$dev" root tbf rate "${rate_kbit}kbit" burst 5k latency 200ms
	done
}

_remove_downstream_limit() {
	local iface dev
	[ -f "${_TSTATE_DIR}/${OMR_QUOTA_INTERFACE}.downstream" ] || return
	for iface in ${OMR_QUOTA_DOWN_INTERFACES:-${OMR_QUOTA_INTERFACES:-$OMR_QUOTA_INTERFACE}}; do
		dev="$(_get_real_interface "$iface")"
		[ -n "$dev" ] && tc qdisc del dev "$dev" root 2>/dev/null || true
	done
	rm -f "${_TSTATE_DIR}/${OMR_QUOTA_INTERFACE}.downstream"
}

_block_lan() {
	[ "${OMR_QUOTA_BLOCK_LAN:-0}" = "1" ] || return
	[ "$(uci -q get firewall.zone_lan.input)" = "DROP" ] && return

	logger -t "OMR-QUOTA" "Block LAN interfaces: $1"
	uci -q set firewall.zone_lan.input='DROP'
	uci -q commit firewall
	/etc/init.d/firewall reload >/dev/null 2>&1
	if [ "$(uci -q get shadowsocks-rust.sss0)" = "ss_rules" ]; then
		uci -q set shadowsocks-rust.sss0.disabled=1
		uci -q commit shadowsocks-rust
		/etc/init.d/shadowsocks-rust stop >/dev/null 2>&1
	fi
}

_unblock_lan() {
	[ "${OMR_QUOTA_BLOCK_LAN:-0}" = "1" ] || return
	[ "$(uci -q get firewall.zone_lan.input)" = "ACCEPT" ] && return

	logger -t "OMR-QUOTA" "Unblock LAN interfaces"
	if [ "$(uci -q get shadowsocks-rust.sss0)" = "ss_rules" ]; then
		uci -q set shadowsocks-rust.sss0.disabled=0
		uci -q commit shadowsocks-rust
		/etc/init.d/shadowsocks-rust start >/dev/null 2>&1
	fi
	uci -q set firewall.zone_lan.input='ACCEPT'
	uci -q commit firewall
	/etc/init.d/firewall reload >/dev/null 2>&1
}

_calculate_budget_limit() {
	local tt="$1"
	local now today midnight end_ts today_ts rd rv dv mm seconds_left percent threshold

	[ -n "$OMR_QUOTA_TT" ] && [ "$OMR_QUOTA_TT" -gt 0 ] || return
	[ -n "$OMR_QUOTA_ENDDATE" ] || return

	percent=$(( tt * 100 / OMR_QUOTA_TT ))
	threshold="${OMR_QUOTA_PERCENT:-80}"
	[ "$percent" -gt "$threshold" ] || return

	end_ts="$(date -d "$OMR_QUOTA_ENDDATE" +%s 2>/dev/null)" || return
	today_ts="$(date -d "00:00" +%s 2>/dev/null)" || return
	rd=$(( (end_ts - today_ts) / 86400 + 1 ))
	[ "$rd" -gt 0 ] || rd=1

	rv=$(( OMR_QUOTA_TT - tt ))
	[ "$rv" -lt 0 ] && rv=0
	dv=$(( rv / rd ))

	now="$(date +%s)"
	today="$(date +%F)"
	midnight="$(date -d "$today 0" +%s 2>/dev/null)" || return
	mm=$(( (now - midnight) / 60 ))
	[ "$mm" -lt 0 ] && mm=0
	[ "$mm" -ge 1439 ] && mm=1439

	if [ "${OMR_QUOTA_METHOD:-0}" = "1" ]; then
		if [ "$(( now - last_calculation ))" -gt "${OMR_QUOTA_CALCULATION_INTERVAL:-120}" ]; then
			cv=$(( dv / (1440 - mm + 1) ))
			last_tt=$tt
			last_calculation=$now
		fi
	elif [ "${OMR_QUOTA_METHOD:-0}" = "2" ]; then
		seconds_left=$(( (1440 - mm) * 60 ))
		[ "$seconds_left" -le 0 ] && seconds_left=60
		cb=$(( dv * 8 / seconds_left ))
		[ "$cb" -lt 1 ] && cb=1
	fi
}

# Baseline reset -- reset_exceeded (UCI/ubus) only clears the *persistent*
# scope marker, which is irrelevant for exceedance_scope=month_only: that
# scope recomputes "exceeded" from live vnstat totals every loop, so there's
# nothing else to clear and quota enforcement stays in effect until the
# calendar month rolls over. Recording a baseline (usage-so-far, subtracted
# from every future reading) makes reset_exceeded actually un-exceed a
# month_only quota immediately, without waiting for vnstat's own month
# bucket to roll over. Tagged with the calendar month it was taken in, so a
# real month rollover (vnstat's own bucket resetting to 0) isn't permanently
# masked by a stale baseline from a previous month.
_baseline_file="${_PERSIST_DIR}/${OMR_QUOTA_INTERFACE}.baseline"

_read_baseline() {
	local tag rx0 tx0
	[ -f "$_baseline_file" ] || { printf '0 0'; return; }
	read -r tag rx0 tx0 < "$_baseline_file"
	if [ "$tag" != "$(date +%Y-%m)" ]; then
		rm -f "$_baseline_file"
		printf '0 0'
	else
		printf '%s %s' "${rx0:-0}" "${tx0:-0}"
	fi
}

if [ "${OMR_QUOTA_RESET_BASELINE:-0}" = "1" ]; then
	_rx0=0; _tx0=0
	for iface in ${OMR_QUOTA_INTERFACES:-$OMR_QUOTA_INTERFACE}; do
		real_iface="$(_get_real_interface "$iface")"
		usage="$(_vnstat_usage "$real_iface")"
		_rx0=$(( _rx0 + $(printf '%s' "$usage" | awk '{print $1}') ))
		_tx0=$(( _tx0 + $(printf '%s' "$usage" | awk '{print $2}') ))
	done
	mkdir -p "$_PERSIST_DIR"
	printf '%s %s %s' "$(date +%Y-%m)" "$_rx0" "$_tx0" > "$_baseline_file"
	logger -t "OMR-QUOTA" "Reset quota baseline for $OMR_QUOTA_INTERFACE (rx=${_rx0}KB tx=${_tx0}KB)"
fi

_prev_exceeded=-1
last_calculation=0
last_tt=0

# main loop
while true; do
	# The interfaces a quota meters are also the interfaces it enforces on:
	# a per-interface quota's list is just itself, a global quota's list is
	# every interface it combines -- so exceeding it cuts/throttles all of them.
	target_interfaces="${OMR_QUOTA_INTERFACES:-$OMR_QUOTA_INTERFACE}"

	rx=0
	tx=0
	for iface in $target_interfaces; do
		real_iface="$(_get_real_interface "$iface")"
		usage="$(_vnstat_usage "$real_iface")"
		rxi="$(printf '%s' "$usage" | awk '{print $1}')"
		txi="$(printf '%s' "$usage" | awk '{print $2}')"
		rx=$(( rx + ${rxi:-0} ))
		tx=$(( tx + ${txi:-0} ))
	done
	_baseline="$(_read_baseline)"
	rx=$(( rx - $(printf '%s' "$_baseline" | awk '{print $1}') ))
	tx=$(( tx - $(printf '%s' "$_baseline" | awk '{print $2}') ))
	[ "$rx" -lt 0 ] && rx=0
	[ "$tx" -lt 0 ] && tx=0
	tt=$(( rx + tx ))
	cb=""

	exceeded=0
	reason=""
	if [ -n "$OMR_QUOTA_RX" ] && [ "$OMR_QUOTA_RX" -gt 0 ] && [ -n "$rx" ] && [ "$OMR_QUOTA_RX" -le "$rx" ]; then
		exceeded=1; reason="RX quota"
	elif [ -n "$OMR_QUOTA_TX" ] && [ "$OMR_QUOTA_TX" -gt 0 ] && [ -n "$tx" ] && [ "$OMR_QUOTA_TX" -le "$tx" ]; then
		exceeded=1; reason="TX quota"
	elif [ -n "$OMR_QUOTA_TT" ] && [ "$OMR_QUOTA_TT" -gt 0 ] && [ -n "$tt" ] && [ "$OMR_QUOTA_TT" -le "$tt" ]; then
		exceeded=1; reason="RX+TX quota"
	fi

	_calculate_budget_limit "$tt"
	if [ -n "$cv" ] && [ "$(( tt - last_tt ))" -gt "$cv" ]; then
		exceeded=1; reason="daily budget"
	fi

	# For persistent scope: once exceeded, stay exceeded even across month boundaries
	_persist_file="${_PERSIST_DIR}/${OMR_QUOTA_INTERFACE}.exceeded"
	if [ "${OMR_QUOTA_SCOPE:-month_only}" = "persistent" ] && [ -f "$_persist_file" ]; then
		exceeded=1
		[ -z "$reason" ] && reason="persistent exceeded state"
	fi

	_tstate_file="${_TSTATE_DIR}/${OMR_QUOTA_INTERFACE}.throttled"

	if [ -n "$cb" ]; then
		_apply_downstream_limit "$cb"
	else
		_remove_downstream_limit
	fi

	if [ "$exceeded" = "1" ]; then
		# Record persistent exceeded state on first detection
		if [ "${OMR_QUOTA_SCOPE:-month_only}" = "persistent" ] && [ ! -f "$_persist_file" ]; then
			mkdir -p "$_PERSIST_DIR"
			touch "$_persist_file"
		fi

		if [ "${OMR_QUOTA_ACTION:-cut}" = "throttle" ]; then
			[ "$_prev_exceeded" != "1" ] && \
				logger -t "OMR-QUOTA" "Throttling $target_interfaces to ${OMR_QUOTA_THROTTLE_DL:-1}/${OMR_QUOTA_THROTTLE_UL:-1} Mbps: $reason reached"
			mkdir -p "$_TSTATE_DIR"
			touch "$_tstate_file"
			for iface in $target_interfaces; do
				real_iface="$(_get_real_interface "$iface")"
				[ -n "$real_iface" ] && _apply_throttle "$real_iface"
				# Keep the interface up while throttled
				iface_up="$(ifstatus "$iface" | jsonfilter -e '@.up')"
				[ "$iface_up" = "false" ] && ifup $iface
			done
		else
			# cut: bring the interface(s) down
			_block_lan "$reason reached"
			for iface in $target_interfaces; do
				iface_up="$(ifstatus "$iface" | jsonfilter -e '@.up')"
				if [ "$iface_up" = "true" ]; then
					logger -t "OMR-QUOTA" "Set interface $iface down: $reason reached"
					ifdown $iface
				fi
			done
		fi
	else
		# Quota not exceeded — remove any throttle that was active
		if [ -f "$_tstate_file" ]; then
			logger -t "OMR-QUOTA" "Removing throttle from $target_interfaces"
			for iface in $target_interfaces; do
				real_iface="$(_get_real_interface "$iface")"
				[ -n "$real_iface" ] && _remove_throttle "$real_iface"
			done
			rm -f "$_tstate_file"
		fi
		for iface in $target_interfaces; do
			iface_up="$(ifstatus "$iface" | jsonfilter -e '@.up')"
			if [ "$iface_up" = "false" ]; then
				logger -t "OMR-QUOTA" "Set interface $iface up"
				ifup $iface
			fi
		done
		_unblock_lan
	fi

	_prev_exceeded=$exceeded
	sleep "${OMR_QUOTA_INTERVAL:-60}"
done
