#!/usr/bin/env bash
set -euo pipefail

read_cfg() {
  python3 - "$@" <<'PY'
import json,sys
cands=["/var/lib/snowpanel/app.json","/etc/snowpanel/app.json","/etc/snowpanel/limits.json","/var/lib/snowpanel/limits.json"]
rate=0
for p in cands:
    try:
        with open(p,"r",encoding="utf-8") as f:
            j=json.load(f)
        if isinstance(j,dict):
            v=j.get("rate_mbps") or (j.get("limits",{}) if isinstance(j.get("limits"),dict) else {}).get("rate_mbps")
            if v is not None:
                rate=int(v)
                break
    except Exception:
        pass
print(rate)
PY
}

get_iface() {
  local ifs=() d4 d6
  d4=$(ip -o -4 route show to default 2>/dev/null | awk '{print $5}' | head -n1 || true)
  d6=$(ip -o -6 route show ::/0      2>/dev/null | awk '{print $5}' | head -n1 || true)
  [[ -n "${d4:-}" ]] && ifs+=("$d4")
  [[ -n "${d6:-}" && "${d6:-}" != "${d4:-}" ]] && ifs+=("$d6")
  [[ ${#ifs[@]} -eq 0 ]] && { echo ""; return 1; }
  echo "${ifs[0]}"
}

get_uid() {
  local u pid
  u=$(systemctl show -p User --value snowflake-proxy 2>/dev/null || true)
  if [[ -n "${u:-}" ]]; then id -u "$u" 2>/dev/null || true; return 0; fi
  if id -u snowflake &>/dev/null; then id -u snowflake; return 0; fi
  pid=$(systemctl show -p MainPID --value snowflake-proxy 2>/dev/null || true)
  if [[ -n "${pid:-}" && -r "/proc/$pid/status" ]]; then awk '/^Uid:/{print $2; exit}' "/proc/$pid/status"; return 0; fi
  echo ""
}

ipt_add() {
  local bin="$1" chain="$2" rule="$3"
  "$bin" -t mangle -N "$chain" 2>/dev/null || true
  "$bin" -t mangle -C OUTPUT -j "$chain" 2>/dev/null || "$bin" -t mangle -A OUTPUT -j "$chain"
  eval "$bin -t mangle -C $chain $rule" 2>/dev/null || eval "$bin -t mangle -A $chain $rule"
}
ipt_in_add() {
  local bin="$1" chain="$2"
  "$bin" -t mangle -N "$chain" 2>/dev/null || true
  "$bin" -t mangle -C PREROUTING -j "$chain" 2>/dev/null || "$bin" -t mangle -A PREROUTING -j "$chain"
  "$bin" -t mangle -C "$chain" -j CONNMARK --restore-mark 2>/dev/null || "$bin" -t mangle -A "$chain" -j CONNMARK --restore-mark
}
ipt_clear() {
  local bin="$1"
  for c in SNOWPANEL SNOWPANEL_IN; do
    while "$bin" -t mangle -D OUTPUT -j SNOWPANEL    2>/dev/null; do :; done
    while "$bin" -t mangle -D PREROUTING -j SNOWPANEL_IN 2>/dev/null; do :; done
    "$bin" -t mangle -F "$c" 2>/dev/null || true
    "$bin" -t mangle -X "$c" 2>/dev/null || true
  done
}

tc_clear() {
  local ifc="$1"
  tc qdisc del dev "$ifc" root    2>/dev/null || true
  tc qdisc del dev "$ifc" ingress 2>/dev/null || true
  tc qdisc del dev ifb0 root      2>/dev/null || true
  tc qdisc del dev ifb0 ingress   2>/dev/null || true
  ip link set ifb0 down           2>/dev/null || true
  ip link delete ifb0 type ifb    2>/dev/null || true
}

tc_apply() {
  local ifc="$1" rate_kbit="$2"
  tc qdisc replace dev "$ifc" root handle 1: htb default 30
  tc class add    dev "$ifc" parent 1: classid 1:1  htb rate 10000000kbit ceil 10000000kbit 2>/dev/null || \
  tc class change dev "$ifc" parent 1: classid 1:1  htb rate 10000000kbit ceil 10000000kbit
  tc class add    dev "$ifc" parent 1:1 classid 1:10 htb rate "${rate_kbit}kbit" ceil "${rate_kbit}kbit" 2>/dev/null || \
  tc class change dev "$ifc" parent 1:1 classid 1:10 htb rate "${rate_kbit}kbit" ceil "${rate_kbit}kbit"
  tc class add    dev "$ifc" parent 1:1 classid 1:30 htb rate 9000000kbit ceil 10000000kbit 2>/dev/null || \
  tc class change dev "$ifc" parent 1:1 classid 1:30 htb rate 9000000kbit ceil 10000000kbit
  tc filter replace dev "$ifc" parent 1: protocol all handle 0x1 fw flowid 1:10

  modprobe ifb numifbs=1 2>/dev/null || true
  ip link add ifb0 type ifb 2>/dev/null || true
  ip link set dev ifb0 up
  tc qdisc replace dev "$ifc" ingress
  tc filter replace dev "$ifc" parent ffff: protocol all u32 match u32 0 0 action mirred egress redirect dev ifb0

  tc qdisc replace dev ifb0 root handle 2: htb default 30
  tc class add    dev ifb0 parent 2: classid 2:1  htb rate 10000000kbit ceil 10000000kbit 2>/dev/null || \
  tc class change dev ifb0 parent 2: classid 2:1  htb rate 10000000kbit ceil 10000000kbit
  tc class add    dev ifb0 parent 2:1 classid 2:10 htb rate "${rate_kbit}kbit" ceil "${rate_kbit}kbit" 2>/dev/null || \
  tc class change dev ifb0 parent 2:1 classid 2:10 htb rate "${rate_kbit}kbit" ceil "${rate_kbit}kbit"
  tc class add    dev ifb0 parent 2:1 classid 2:30 htb rate 9000000kbit ceil 10000000kbit 2>/dev/null || \
  tc class change dev ifb0 parent 2:1 classid 2:30 htb rate 9000000kbit ceil 10000000kbit
  tc filter replace dev ifb0 parent 2: protocol all handle 0x1 fw flowid 2:10
}

apply() {
  local rate uid ifc per_kbit
  rate=$(read_cfg)
  uid=$(get_uid)
  ifc=$(get_iface || true)
  ipt_clear iptables || true
  ipt_clear ip6tables || true
  if [[ -z "${ifc:-}" || -z "${uid:-}" || "$rate" -le 0 ]]; then
    tc_clear "${ifc:-eth0}" || true
    exit 0
  fi
  per_kbit=$(( rate * 1000 / 2 ))
  [[ $per_kbit -lt 64 ]] && per_kbit=64

  ipt_add iptables  SNOWPANEL   "-m owner --uid-owner $uid -j MARK --set-xmark 0x1/0x1"
  ipt_add iptables  SNOWPANEL   "-m owner --uid-owner $uid -j CONNMARK --save-mark"
  ipt_in_add iptables SNOWPANEL_IN
  if command -v ip6tables >/dev/null 2>&1; then
    ipt_add ip6tables SNOWPANEL   "-m owner --uid-owner $uid -j MARK --set-xmark 0x1/0x1"
    ipt_add ip6tables SNOWPANEL   "-m owner --uid-owner $uid -j CONNMARK --save-mark"
    ipt_in_add ip6tables SNOWPANEL_IN
  fi

  tc_clear "$ifc" || true
  tc_apply "$ifc" "$per_kbit"
}

clear_all() {
  local ifc
  ifc=$(get_iface || echo eth0)
  ipt_clear iptables  || true
  ipt_clear ip6tables || true
  tc_clear "$ifc"     || true
}

cmd="${1:-apply}"
case "$cmd" in
  apply) apply ;;
  clear) clear_all ;;
  *) apply ;;
esac