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

APP_COMPONENT="${APP_COMPONENT:-combined}"
TPROXY_PORT="${TPROXY_PORT:-7895}"
TPROXY_MARK="${TPROXY_MARK:-1}"
TPROXY_TABLE="${TPROXY_TABLE:-100}"
TPROXY_CHAIN="${TPROXY_CHAIN:-VPN_PROXY_TPROXY}"
DEVICE_POLICY_CHAIN="${DEVICE_POLICY_CHAIN:-VPN_PROXY_DEVICE_POLICY}"
GATEWAY_FORWARD_CHAIN="${GATEWAY_FORWARD_CHAIN:-VPN_PROXY_FORWARD}"
GATEWAY_NAT_CHAIN="${GATEWAY_NAT_CHAIN:-VPN_PROXY_NAT}"
TRAFFIC_UPLOAD_CHAIN="${TRAFFIC_UPLOAD_CHAIN:-VPN_PROXY_TRAFFIC_UP}"
TRAFFIC_DOWNLOAD_CHAIN="${TRAFFIC_DOWNLOAD_CHAIN:-VPN_PROXY_TRAFFIC_DOWN}"
DIRECT_TRAFFIC_CHAIN="${DIRECT_TRAFFIC_CHAIN:-VPN_PROXY_DIRECT}"
DIRECT_TRAFFIC_MARK="${DIRECT_TRAFFIC_MARK:-0x40000000}"
GATEWAY_CLIENT_CIDRS="${GATEWAY_CLIENT_CIDRS:-10.0.0.0/8 172.16.0.0/12 192.168.0.0/16}"
PROXY_PORT="${PROXY_PORT:-8080}"
PROXY_BIND_IP="${PROXY_BIND_IP:-0.0.0.0}"
PROXY_INPUT_CHAIN="${PROXY_INPUT_CHAIN:-VPN_PROXY_INPUT}"
PROXY_FIREWALL="${PROXY_FIREWALL:-true}"
PROXY_ALLOWED_CIDRS="${PROXY_ALLOWED_CIDRS:-10.0.0.0/8 172.16.0.0/12 192.168.0.0/16}"
BYPASS_CIDRS="${BYPASS_CIDRS:-0.0.0.0/8 10.0.0.0/8 100.64.0.0/10 127.0.0.0/8 169.254.0.0/16 172.16.0.0/12 192.168.0.0/16 224.0.0.0/4 240.0.0.0/4}"
export TPROXY_PORT TPROXY_MARK DEVICE_POLICY_CHAIN TRAFFIC_UPLOAD_CHAIN TRAFFIC_DOWNLOAD_CHAIN DIRECT_TRAFFIC_CHAIN DIRECT_TRAFFIC_MARK GATEWAY_CLIENT_CIDRS BYPASS_CIDRS
DEVICE_TRAFFIC_CONFIG_VALID=false
export DEVICE_TRAFFIC_ACCOUNTING_ENABLED=false

log() {
  printf '[gateway-entrypoint] %s\n' "$*"
}

if [[ "$APP_COMPONENT" == "control" ]]; then
  exec node /app/dist/server/main.js
fi

validate_device_traffic_config() {
  if [[ -z "$DIRECT_TRAFFIC_CHAIN" || ${#DIRECT_TRAFFIC_CHAIN} -gt 24
    || "$DIRECT_TRAFFIC_CHAIN" =~ [^a-zA-Z0-9_] ]]; then
    log "device traffic counters unavailable: invalid DIRECT_TRAFFIC_CHAIN"
    return 1
  fi
  local direct_names=("$DIRECT_TRAFFIC_CHAIN" "${DIRECT_TRAFFIC_CHAIN}_A" "${DIRECT_TRAFFIC_CHAIN}_B")
  local reserved_names=(
    PREROUTING INPUT FORWARD OUTPUT POSTROUTING
    "$TPROXY_CHAIN"
    "$DEVICE_POLICY_CHAIN" "${DEVICE_POLICY_CHAIN}_A" "${DEVICE_POLICY_CHAIN}_B"
    "$TRAFFIC_DOWNLOAD_CHAIN" "${TRAFFIC_DOWNLOAD_CHAIN}_A" "${TRAFFIC_DOWNLOAD_CHAIN}_B"
    "${TRAFFIC_DOWNLOAD_CHAIN}_A_P" "${TRAFFIC_DOWNLOAD_CHAIN}_B_P"
  )
  for direct_name in "${direct_names[@]}"; do
    for reserved_name in "${reserved_names[@]}"; do
      if [[ "$direct_name" == "$reserved_name" ]]; then
        log "device traffic counters unavailable: DIRECT_TRAFFIC_CHAIN conflicts with ${reserved_name}"
        return 1
      fi
    done
  done
  if ! [[ "$DIRECT_TRAFFIC_MARK" =~ ^(0[xX][0-9a-fA-F]{1,8}|[0-9]{1,10})$
    && "$TPROXY_MARK" =~ ^(0[xX][0-9a-fA-F]{1,8}|[0-9]{1,10})$ ]]; then
    log "device traffic counters unavailable: invalid traffic mark"
    return 1
  fi
  local direct_mark_value tproxy_mark_value
  if [[ "$DIRECT_TRAFFIC_MARK" =~ ^0[xX] ]]; then
    direct_mark_value=$((16#${DIRECT_TRAFFIC_MARK:2}))
  else
    direct_mark_value=$((10#$DIRECT_TRAFFIC_MARK))
  fi
  if [[ "$TPROXY_MARK" =~ ^0[xX] ]]; then
    tproxy_mark_value=$((16#${TPROXY_MARK:2}))
  else
    tproxy_mark_value=$((10#$TPROXY_MARK))
  fi
  if (( direct_mark_value == 0 || direct_mark_value > 0xffffffff
    || (direct_mark_value & (direct_mark_value - 1)) != 0
    || (direct_mark_value & tproxy_mark_value) != 0 )); then
    log "device traffic counters unavailable: DIRECT_TRAFFIC_MARK must be one bit outside TPROXY_MARK"
    return 1
  fi
  DEVICE_TRAFFIC_CONFIG_VALID=true
}

ipt() {
  iptables -w "$@"
}

ipt_traffic() {
  iptables -w 1 "$@"
}

cleanup_proxy_firewall() {
  ipt -D INPUT -p tcp --dport "$PROXY_PORT" -j "$PROXY_INPUT_CHAIN" 2>/dev/null || true
  ipt -D INPUT -p udp --dport "$PROXY_PORT" -j "$PROXY_INPUT_CHAIN" 2>/dev/null || true
  ipt -F "$PROXY_INPUT_CHAIN" 2>/dev/null || true
  ipt -X "$PROXY_INPUT_CHAIN" 2>/dev/null || true
}

cleanup_tproxy() {
  ipt -t mangle -D PREROUTING -j "$TPROXY_CHAIN" 2>/dev/null || true
  ipt -t mangle -F "$TPROXY_CHAIN" 2>/dev/null || true
  ipt -t mangle -X "$TPROXY_CHAIN" 2>/dev/null || true
  ip rule del fwmark "$TPROXY_MARK" table "$TPROXY_TABLE" 2>/dev/null || true
  ip route flush table "$TPROXY_TABLE" 2>/dev/null || true
}

cleanup_device_policy() {
  ipt -t mangle -F "$DEVICE_POLICY_CHAIN" 2>/dev/null || true
  for slot in A B; do
    ipt -t mangle -F "${DEVICE_POLICY_CHAIN}_${slot}" 2>/dev/null || true
    ipt -t mangle -X "${DEVICE_POLICY_CHAIN}_${slot}" 2>/dev/null || true
  done
  ipt -t mangle -X "$DEVICE_POLICY_CHAIN" 2>/dev/null || true
}

setup_device_policy() {
  cleanup_device_policy
  ipt -t mangle -N "$DEVICE_POLICY_CHAIN" || return 1
  ipt -t mangle -N "${DEVICE_POLICY_CHAIN}_A" || return 1
  ipt -t mangle -N "${DEVICE_POLICY_CHAIN}_B" || return 1
  ipt -t mangle -A "${DEVICE_POLICY_CHAIN}_A" -p tcp -j TPROXY --on-port "$TPROXY_PORT" --tproxy-mark "$TPROXY_MARK/$TPROXY_MARK" || return 1
  ipt -t mangle -A "${DEVICE_POLICY_CHAIN}_A" -p udp -j TPROXY --on-port "$TPROXY_PORT" --tproxy-mark "$TPROXY_MARK/$TPROXY_MARK" || return 1
  ipt -t mangle -A "$DEVICE_POLICY_CHAIN" -j "${DEVICE_POLICY_CHAIN}_A" || return 1
}

cleanup_gateway_forwarding() {
  ipt -D FORWARD -j "$GATEWAY_FORWARD_CHAIN" 2>/dev/null || true
  ipt -t nat -D POSTROUTING -j "$GATEWAY_NAT_CHAIN" 2>/dev/null || true
  ipt -F "$GATEWAY_FORWARD_CHAIN" 2>/dev/null || true
  ipt -X "$GATEWAY_FORWARD_CHAIN" 2>/dev/null || true
  ipt -t nat -F "$GATEWAY_NAT_CHAIN" 2>/dev/null || true
  ipt -t nat -X "$GATEWAY_NAT_CHAIN" 2>/dev/null || true
}

cleanup_device_traffic() {
  [[ "$DEVICE_TRAFFIC_CONFIG_VALID" == "true" ]] || return 0
  ipt_traffic -t mangle -D "$TPROXY_CHAIN" -j CONNMARK --set-xmark "0x0/$DIRECT_TRAFFIC_MARK" 2>/dev/null || true
  ipt_traffic -t raw -D PREROUTING -j "$TRAFFIC_UPLOAD_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -D PREROUTING -j "$DIRECT_TRAFFIC_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -D POSTROUTING -j "$TRAFFIC_DOWNLOAD_CHAIN" 2>/dev/null || true
  ipt_traffic -t raw -F "$TRAFFIC_UPLOAD_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -F "$DIRECT_TRAFFIC_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -F "$TRAFFIC_DOWNLOAD_CHAIN" 2>/dev/null || true
  for slot in A B; do
    ipt_traffic -t raw -F "${TRAFFIC_UPLOAD_CHAIN}_${slot}" 2>/dev/null || true
    ipt_traffic -t raw -F "${TRAFFIC_UPLOAD_CHAIN}_${slot}_P" 2>/dev/null || true
    ipt_traffic -t raw -X "${TRAFFIC_UPLOAD_CHAIN}_${slot}_P" 2>/dev/null || true
    ipt_traffic -t raw -X "${TRAFFIC_UPLOAD_CHAIN}_${slot}" 2>/dev/null || true
    ipt_traffic -t mangle -F "${DIRECT_TRAFFIC_CHAIN}_${slot}" 2>/dev/null || true
    ipt_traffic -t mangle -X "${DIRECT_TRAFFIC_CHAIN}_${slot}" 2>/dev/null || true
    ipt_traffic -t mangle -F "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}" 2>/dev/null || true
    ipt_traffic -t mangle -F "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}_P" 2>/dev/null || true
    ipt_traffic -t mangle -X "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}_P" 2>/dev/null || true
    ipt_traffic -t mangle -X "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}" 2>/dev/null || true
  done
  ipt_traffic -t raw -X "$TRAFFIC_UPLOAD_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -X "$DIRECT_TRAFFIC_CHAIN" 2>/dev/null || true
  ipt_traffic -t mangle -X "$TRAFFIC_DOWNLOAD_CHAIN" 2>/dev/null || true
}

setup_device_traffic() {
  log "setup device traffic counters"
  cleanup_device_traffic
  ipt_traffic -t raw -N "$TRAFFIC_UPLOAD_CHAIN" || return 1
  ipt_traffic -t mangle -N "$DIRECT_TRAFFIC_CHAIN" || return 1
  ipt_traffic -t mangle -N "$TRAFFIC_DOWNLOAD_CHAIN" || return 1
  for slot in A B; do
    ipt_traffic -t raw -N "${TRAFFIC_UPLOAD_CHAIN}_${slot}" || return 1
    ipt_traffic -t raw -N "${TRAFFIC_UPLOAD_CHAIN}_${slot}_P" || return 1
    ipt_traffic -t mangle -N "${DIRECT_TRAFFIC_CHAIN}_${slot}" || return 1
    ipt_traffic -t mangle -N "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}" || return 1
    ipt_traffic -t mangle -N "${TRAFFIC_DOWNLOAD_CHAIN}_${slot}_P" || return 1
  done
  ipt_traffic -t raw -I PREROUTING 1 -j "$TRAFFIC_UPLOAD_CHAIN" || return 1
  # sing-box inserts TPROXY at position 1 later; this jump then sees only packets not intercepted by it.
  ipt_traffic -t mangle -I PREROUTING 1 -j "$DIRECT_TRAFFIC_CHAIN" || return 1
  ipt_traffic -t mangle -I POSTROUTING 1 -j "$TRAFFIC_DOWNLOAD_CHAIN" || return 1
  local policy_rule=4
  for _cidr in $BYPASS_CIDRS; do
    policy_rule=$((policy_rule + 1))
  done
  ipt_traffic -t mangle -I "$TPROXY_CHAIN" "$policy_rule" -j CONNMARK --set-xmark "0x0/$DIRECT_TRAFFIC_MARK" || return 1
}

enable_ip_forwarding() {
  if [[ -w /proc/sys/net/ipv4/ip_forward ]]; then
    printf '1' > /proc/sys/net/ipv4/ip_forward || true
  elif command -v sysctl >/dev/null 2>&1; then
    sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1 || true
  fi
}

setup_proxy_firewall() {
  if [[ "$PROXY_FIREWALL" != "true" || "$PROXY_BIND_IP" == "127.0.0.1" || "$PROXY_BIND_IP" == "::1" ]]; then
    return
  fi

  cleanup_proxy_firewall
  ipt -N "$PROXY_INPUT_CHAIN"
  for cidr in $PROXY_ALLOWED_CIDRS; do
    ipt -A "$PROXY_INPUT_CHAIN" -s "$cidr" -j RETURN
  done
  ipt -A "$PROXY_INPUT_CHAIN" -j DROP
  ipt -I INPUT -p tcp --dport "$PROXY_PORT" -j "$PROXY_INPUT_CHAIN"
  ipt -I INPUT -p udp --dport "$PROXY_PORT" -j "$PROXY_INPUT_CHAIN"
}

setup_gateway_forwarding() {
  log "setup direct gateway forwarding"
  cleanup_gateway_forwarding
  enable_ip_forwarding

  ipt -N "$GATEWAY_FORWARD_CHAIN"
  ipt -t nat -N "$GATEWAY_NAT_CHAIN"
  for cidr in $GATEWAY_CLIENT_CIDRS; do
    ipt -A "$GATEWAY_FORWARD_CHAIN" -s "$cidr" -j ACCEPT
    ipt -A "$GATEWAY_FORWARD_CHAIN" -d "$cidr" -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
    ipt -t nat -A "$GATEWAY_NAT_CHAIN" -s "$cidr" -m addrtype ! --dst-type LOCAL -j MASQUERADE
  done
  ipt -I FORWARD 1 -j "$GATEWAY_FORWARD_CHAIN"
  ipt -t nat -I POSTROUTING 1 -j "$GATEWAY_NAT_CHAIN"
}

setup_tproxy() {
  log "setup tproxy on port ${TPROXY_PORT}"
  cleanup_tproxy
  if ! setup_device_policy; then
    log "device policy unavailable; using global VPN fallback"
    cleanup_device_policy
  fi
  enable_ip_forwarding

  ip rule add fwmark "$TPROXY_MARK" table "$TPROXY_TABLE" 2>/dev/null || true
  ip route replace local 0.0.0.0/0 dev lo table "$TPROXY_TABLE"
  ipt -t mangle -N "$TPROXY_CHAIN"
  ipt -t mangle -A "$TPROXY_CHAIN" -m addrtype --dst-type LOCAL -j RETURN
  ipt -t mangle -A "$TPROXY_CHAIN" -m mark --mark "$TPROXY_MARK" -j RETURN
  ipt -t mangle -A "$TPROXY_CHAIN" -i 'br-+' -j RETURN

  # Private/local destinations stay reachable; every intercepted public packet goes to VPN.
  for cidr in $BYPASS_CIDRS; do
    ipt -t mangle -A "$TPROXY_CHAIN" -d "$cidr" -j RETURN
  done
  if ipt -t mangle -L "$DEVICE_POLICY_CHAIN" -n >/dev/null 2>&1; then
    ipt -t mangle -A "$TPROXY_CHAIN" -j "$DEVICE_POLICY_CHAIN"
  else
    ipt -t mangle -A "$TPROXY_CHAIN" -p tcp -j TPROXY --on-port "$TPROXY_PORT" --tproxy-mark "$TPROXY_MARK/$TPROXY_MARK"
    ipt -t mangle -A "$TPROXY_CHAIN" -p udp -j TPROXY --on-port "$TPROXY_PORT" --tproxy-mark "$TPROXY_MARK/$TPROXY_MARK"
  fi
}

setup_gateway_forwarding
setup_tproxy
if validate_device_traffic_config; then
  if ! setup_device_traffic; then
    log "device traffic counters unavailable; VPN routing remains active"
    cleanup_device_traffic
  else
    export DEVICE_TRAFFIC_ACCOUNTING_ENABLED=true
  fi
fi
setup_proxy_firewall

node /app/dist/server/main.js &
APP_PID=$!

shutdown() {
  kill "$APP_PID" 2>/dev/null || true
  wait "$APP_PID" 2>/dev/null || true
  cleanup_proxy_firewall
  cleanup_device_traffic
  cleanup_tproxy
  cleanup_device_policy
  cleanup_gateway_forwarding
}

trap 'shutdown; exit 0' SIGTERM SIGINT
wait "$APP_PID"
STATUS=$?
cleanup_proxy_firewall
cleanup_device_traffic
cleanup_tproxy
cleanup_device_policy
cleanup_gateway_forwarding
exit "$STATUS"