#!/bin/bash
set -euo pipefail

# ----------------------------
# Config
# ----------------------------
MIHOMO_UID="mihomo"
REDIR_PORT="7892"     # TCP Redirect
TPROXY_PORT="7893"    # UDP/TCP TProxy
FW_MARK="0x1"
ROUTE_TABLE="100"

EXCLUDE_IFACES=("tun0" "wg0")

# ----------------------------
# Helpers
# ----------------------------
ipt() { iptables "$@"; }

# Надежная функция удаления: Пытается удалить правило, пока iptables не скажет "нет такого правила"
# $1=table, $2=chain, $3=args (часть правила для матчинга)
del_loop() {
  local table=$1
  local chain=$2
  shift 2
  local rule_args="$@"

  # Пока проверка (-C) успешна, делаем удаление (-D)
  while iptables -t "$table" -C "$chain" $rule_args 2>/dev/null; do
    echo "Deleting from $table/$chain: $rule_args"
    iptables -t "$table" -D "$chain" $rule_args
  done
}

ensure_ip_rule() {
  while ip rule list | grep -q "fwmark ${FW_MARK} lookup ${ROUTE_TABLE}"; do
    ip rule del fwmark ${FW_MARK} lookup ${ROUTE_TABLE} || true
  done
  ip rule add fwmark ${FW_MARK} lookup ${ROUTE_TABLE}
  ip route replace local 0.0.0.0/0 dev lo table ${ROUTE_TABLE}
}

# ----------------------------
# CLEANUP PHASE
# ----------------------------
echo "--- Cleaning up old rules (Robust Mode) ---"

# 1. Удаляем ссылки (JUMP) на наши цепочки
del_loop nat OUTPUT -p tcp -m comment --comment "MIHOMO-JUMP" -j MIHOMO_REDIR
del_loop nat PREROUTING -i wt0 -p tcp -m comment --comment "MIHOMO-REDIRECT" -j REDIRECT --to-port "${REDIR_PORT}"
del_loop mangle PREROUTING -i wt0 -m comment --comment "MIHOMO-JUMP" -j MIHOMO_TPROXY

# 2. Удаляем исключения (EXCLUDE) и маркеры (MARK)
# Mangle OUTPUT
del_loop mangle OUTPUT -p tcp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"
del_loop mangle OUTPUT -p udp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"

for IFACE in "${EXCLUDE_IFACES[@]}"; do
  del_loop mangle OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
  del_loop mangle PREROUTING -i "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
  del_loop nat OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

del_loop mangle OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
del_loop nat OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

# 3. Теперь, когда ссылок нет, можно безопасно удалить цепочки
# Сначала сбрасываем содержимое (-F), потом удаляем саму цепочку (-X)
ipt -t mangle -F MIHOMO_TPROXY 2>/dev/null || true
ipt -t mangle -X MIHOMO_TPROXY 2>/dev/null || true

ipt -t nat -F MIHOMO_REDIR 2>/dev/null || true
ipt -t nat -X MIHOMO_REDIR 2>/dev/null || true

echo "--- Cleanup finished. Applying new rules ---"

# ----------------------------
# NAT (REDIRECT) - TCP
# ----------------------------
ipt -t nat -N MIHOMO_REDIR

# Local exclusions inside chain
ipt -t nat -A MIHOMO_REDIR -d 192.168.0.0/16 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 10.0.0.0/8 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 172.16.0.0/12 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 127.0.0.0/8 -j RETURN
ipt -t nat -A MIHOMO_REDIR -p tcp -j REDIRECT --to-ports "${REDIR_PORT}"

# Apply to OUTPUT (Local)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t nat -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done
ipt -t nat -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
ipt -t nat -A OUTPUT -p tcp -m comment --comment "MIHOMO-JUMP" -j MIHOMO_REDIR

# Apply to PREROUTING (wt0 Ingress) - Force Redir
ipt -t nat -A PREROUTING -i wt0 -p tcp -m comment --comment "MIHOMO-REDIRECT" -j REDIRECT --to-port "${REDIR_PORT}"

# ----------------------------
# MANGLE (TPROXY) - UDP/TCP
# ----------------------------
ensure_ip_rule
ipt -t mangle -N MIHOMO_TPROXY

# Local exclusions inside chain
ipt -t mangle -A MIHOMO_TPROXY -d 192.168.0.0/16 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 10.0.0.0/8 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 172.16.0.0/12 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 127.0.0.0/8 -j RETURN

# TProxy Targets
# ipt -t mangle -A MIHOMO_TPROXY -p tcp -j TPROXY --on-port "${TPROXY_PORT}" --tproxy-mark "${FW_MARK}/${FW_MARK}"
ipt -t mangle -A MIHOMO_TPROXY -p udp -j TPROXY --on-port "${TPROXY_PORT}" --tproxy-mark "${FW_MARK}/${FW_MARK}"

# Apply to OUTPUT (Local)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done
ipt -t mangle -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

# Mark packets
# ipt -t mangle -A OUTPUT -p tcp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"
ipt -t mangle -A OUTPUT -p udp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"

# Apply to PREROUTING (wt0 Ingress)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A PREROUTING -i "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done
ipt -t mangle -A PREROUTING -i wt0 -m comment --comment "MIHOMO-JUMP" -j MIHOMO_TPROXY

echo "Done."