#!/bin/bash

# --- Script location & config ---
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
CONFIG_FILE="$SCRIPT_DIR/lxc_bandwidth_limit.conf"

# --- UDM WAN interface ---
UDM_IP="192.168.1.1"
UDM_USER="root"
UDM_IF="ppp0"

# --- Default values (overridden by config/CLI) ---
CHECK_INTERVAL=60
DRY_RUN=0
LIMIT_APPLIED=0

# --- Load Apprise URL ---
[[ -f /root/apprise.env ]] && source /root/apprise.env

# --- Parse config file ---
if [[ -f "$CONFIG_FILE" ]]; then
  echo "📄 Loading config from $CONFIG_FILE"
  while read -r line; do
    [[ "$line" =~ ^#.*$ || -z "$line" ]] && continue
    if [[ -z "${HEADERS[*]}" ]]; then
      IFS='|' read -ra HEADERS <<< "$line"
      for i in "${!HEADERS[@]}"; do
        HEADERS[$i]=$(echo "${HEADERS[$i]}" | xargs | tr '-' '_')
      done
    else
      IFS='|' read -ra VALUES <<< "$line"
      for i in "${!VALUES[@]}"; do
        val=$(echo "${VALUES[$i]}" | xargs)
        key="${HEADERS[$i]}"
        export "$key=$val"
      done
      break
    fi
  done < "$CONFIG_FILE"
fi

# --- Parse CLI arguments ---
for arg in "$@"; do
  case $arg in
    --config=*) CONFIG_FILE="${arg#*=}" ;;
    --container_id=*) container_id="${arg#*=}" ;;
    --capacity-up=*) capacity_up="${arg#*=}" ;;
    --capacity-down=*) capacity_down="${arg#*=}" ;;
    --threshold-up=*) threshold_up="${arg#*=}" ;;
    --threshold-down=*) threshold_down="${arg#*=}" ;;
    --activate-limit=*) activate_limit="${arg#*=}" ;;
    --dry-run) DRY_RUN=1 ;;
    --status)
      read -r RX1 TX1 < <(ssh -o BatchMode=yes "$UDM_USER@$UDM_IP" "awk '/$UDM_IF/ {print \$2, \$10}' /proc/net/dev")
      sleep 1
      read -r RX2 TX2 < <(ssh -o BatchMode=yes "$UDM_USER@$UDM_IP" "awk '/$UDM_IF/ {print \$2, \$10}' /proc/net/dev")
      RX_MBPS=$(( (RX2 - RX1) * 8 / 1000000 ))
      TX_MBPS=$(( (TX2 - TX1) * 8 / 1000000 ))
      echo "📥 RX: ${RX_MBPS} Mbps | 📤 TX: ${TX_MBPS} Mbps"
      echo "🔧 Config:"
      echo "  Container ID     : ${container_id}"
      echo "  Capacity Up      : ${capacity_up} Mbps"
      echo "  Capacity Down    : ${capacity_down} Mbps"
      echo "  Threshold Up     : ${threshold_up}% → $((capacity_up * threshold_up / 100)) Mbps"
      echo "  Threshold Down   : ${threshold_down}% → $((capacity_down * threshold_down / 100)) Mbps"
      echo "  Activate Limit   : ${activate_limit}"
      if [[ "$activate_limit" == "on" ]] && \
         (( TX_MBPS > (capacity_up * threshold_up / 100) || RX_MBPS > (capacity_down * threshold_down / 100) )); then
        echo "🚨 Limit WOULD be applied."
      else
        echo "✅ Limit would not be applied."
      fi
      exit 0
      ;;
    --help)
      echo "Usage: $0 [--dry-run] [--status] [--config=...] [--container_id=...] ..."
      exit 0
      ;;
    *) echo "❌ Unknown parameter: $arg" && exit 1 ;;
  esac
done

# --- Validation ---
if [[ -z "${container_id:-}" ]]; then
  echo "❌ Missing required: container_id"
  exit 1
fi

# --- Functions ---
notify_apprise() {
  [[ -n "${APPRISE_URL:-}" && $DRY_RUN -eq 0 ]] && apprise -b "$1" "$APPRISE_URL" &
}

get_lxc_veth_if() {
  grep -l "$container_id" /sys/class/net/*/iflink | cut -d/ -f5
}

apply_limit() {
  local IFACE="$1"
  local LIMIT=$(($capacity_up < $capacity_down ? $capacity_up : $capacity_down))
  if [[ $DRY_RUN -eq 1 ]]; then
    echo "💡 [Dry-run] Would apply limit of ${LIMIT}Mbps to $IFACE"
  else
    tc qdisc replace dev "$IFACE" root tbf rate "$((LIMIT * 125000))bps" burst 32kbit latency 400ms
    echo "⚠️ Applied limit of ${LIMIT}Mbps to $IFACE"
    notify_apprise "⚠️ Limit applied to CT $container_id: ${LIMIT}Mbps"
  fi
  LIMIT_APPLIED=1
}

remove_limit() {
  local IFACE="$1"
  if [[ $DRY_RUN -eq 1 ]]; then
    echo "💡 [Dry-run] Would remove limit from $IFACE"
  else
    tc qdisc del dev "$IFACE" root 2>/dev/null || true
    echo "✅ Removed limit from $IFACE"
    notify_apprise "✅ Limit removed from CT $container_id"
  fi
  LIMIT_APPLIED=0
}

# --- Main loop ---
while true; do
  read -r RX1 TX1 < <(ssh -o BatchMode=yes "$UDM_USER@$UDM_IP" "awk '/$UDM_IF/ {print \$2, \$10}' /proc/net/dev")
  sleep 1
  read -r RX2 TX2 < <(ssh -o BatchMode=yes "$UDM_USER@$UDM_IP" "awk '/$UDM_IF/ {print \$2, \$10}' /proc/net/dev")
  RX_MBPS=$(( (RX2 - RX1) * 8 / 1000000 ))
  TX_MBPS=$(( (TX2 - TX1) * 8 / 1000000 ))
  echo "📥 RX: ${RX_MBPS} Mbps | 📤 TX: ${TX_MBPS} Mbps"

  IFACE=$(get_lxc_veth_if)
  [[ -z "$IFACE" ]] && echo "❌ Could not find veth for CT $container_id" && sleep $CHECK_INTERVAL && continue

  RX_TRIGGER=$((capacity_down * threshold_down / 100))
  TX_TRIGGER=$((capacity_up * threshold_up / 100))

  if [[ "$activate_limit" == "on" ]]; then
    if (( RX_MBPS > RX_TRIGGER || TX_MBPS > TX_TRIGGER )); then
      [[ $LIMIT_APPLIED -eq 0 ]] && apply_limit "$IFACE"
    else
      [[ $LIMIT_APPLIED -eq 1 ]] && remove_limit "$IFACE"
    fi
  else
    echo "🚫 Limit logic disabled (activate-limit=off)"
    [[ $LIMIT_APPLIED -eq 1 ]] && remove_limit "$IFACE"
  fi

  sleep $CHECK_INTERVAL
done