#!/bin/bash
set -euo pipefail

# ----------------------------
# Config
# ----------------------------
MIHOMO_UID="mihomo"

REDIR_PORT="7892"     # mihomo redir-port (NAT REDIRECT for TCP)
TPROXY_PORT="7893"    # mihomo tproxy-port (TPROXY for TCP/UDP)
FW_MARK="0x1"
ROUTE_TABLE="100"

# Interfaces to EXCLUDE completely from interception
EXCLUDE_IFACES=("tun0" "wg0")

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

ensure_ip_rule() {
  # Route marked traffic to local via custom table (idempotent-ish)
  ip rule list | grep -q "fwmark ${FW_MARK} lookup ${ROUTE_TABLE}" || \
    ip rule add fwmark ${FW_MARK} lookup ${ROUTE_TABLE}

  # Route everything in that table to local loopback so TPROXY can catch it
  ip route show table ${ROUTE_TABLE} | grep -q "^local 0.0.0.0/0 dev lo" || \
    ip route add local 0.0.0.0/0 dev lo table ${ROUTE_TABLE}
}

# ----------------------------
# NAT (REDIRECT) part (TCP only)
# ----------------------------

# Cleanup old chains (ignore if absent)
ipt -t nat -F MIHOMO_REDIR 2>/dev/null || true
ipt -t nat -X MIHOMO_REDIR 2>/dev/null || true

# NOTE: Your original script flushes OUTPUT table nat globally.
# Keeping behavior to match your current approach, but yes, it nukes other rules.
ipt -t nat -F OUTPUT 2>/dev/null || true

# Create chain
ipt -t nat -N MIHOMO_REDIR

# Exclude loopback and local subnets
ipt -t nat -A MIHOMO_REDIR -d 127.0.0.0/8 -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 192.168.0.0/16 -j RETURN

# Exclude traffic that is going via tun0/wg0 (your "do not touch" tunnels)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t nat -A OUTPUT -o "${IFACE}" -j RETURN 2>/dev/null || true
done

# Exclude mihomo's own traffic by UID
ipt -t nat -C OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -j RETURN 2>/dev/null || \
ipt -t nat -I OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -j RETURN

# Everything else TCP -> REDIRECT to mihomo
ipt -t nat -A MIHOMO_REDIR -p tcp -j REDIRECT --to-ports "${REDIR_PORT}"

# Apply to local OUTPUT TCP
ipt -t nat -C OUTPUT -p tcp -j MIHOMO_REDIR 2>/dev/null || \
ipt -t nat -A OUTPUT -p tcp -j MIHOMO_REDIR

# Apply to traffic coming from NetBird interface (ingress)
# Exclude tun0/wg0 by design: only target wt0 here.
ipt -t nat -C PREROUTING -i wt0 -p tcp -j REDIRECT --to-port "${REDIR_PORT}" 2>/dev/null || \
ipt -t nat -A PREROUTING -i wt0 -p tcp -j REDIRECT --to-port "${REDIR_PORT}"

# IMPORTANT:
# Removed your old rule "PREROUTING -i wg0 ... REDIRECT"
# because you explicitly asked to exclude wg0 from routing/interception.

# ----------------------------
# MANGLE (TPROXY) part (TCP+UDP typically)
# ----------------------------

ensure_ip_rule

# Cleanup old chains
ipt -t mangle -F MIHOMO_TPROXY 2>/dev/null || true
ipt -t mangle -X MIHOMO_TPROXY 2>/dev/null || true

# Create chain
ipt -t mangle -N MIHOMO_TPROXY

# Exclusions: loopback/local subnets
ipt -t mangle -A MIHOMO_TPROXY -d 127.0.0.0/8 -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 192.168.0.0/16 -j RETURN

# Exclude traffic arriving from tun0/wg0 (ingress side)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A PREROUTING -i "${IFACE}" -j RETURN 2>/dev/null || true
done

# Exclude traffic leaving via tun0/wg0 (local OUTPUT side)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -C OUTPUT -o "${IFACE}" -j RETURN 2>/dev/null || \
  ipt -t mangle -I OUTPUT -o "${IFACE}" -j RETURN
done

# Exclude mihomo's own traffic (OUTPUT) so it doesn't eat itself
ipt -t mangle -C OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -j RETURN 2>/dev/null || \
ipt -t mangle -I OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -j RETURN

# --- TPROXY for local VM traffic (OUTPUT) ---
# Mark TCP/UDP so policy routing sends them to lo (where mihomo tproxy listens)
ipt -t mangle -C OUTPUT -p tcp -j MARK --set-mark "${FW_MARK}" 2>/dev/null || \
ipt -t mangle -A OUTPUT -p tcp -j MARK --set-mark "${FW_MARK}"

ipt -t mangle -C OUTPUT -p udp -j MARK --set-mark "${FW_MARK}" 2>/dev/null || \
ipt -t mangle -A OUTPUT -p udp -j MARK --set-mark "${FW_MARK}"

# --- TPROXY for wt0 ingress traffic (PREROUTING) ---
# First run through our exclusions, then TPROXY it.
ipt -t mangle -C PREROUTING -i wt0 -j MIHOMO_TPROXY 2>/dev/null || \
ipt -t mangle -A PREROUTING -i wt0 -j MIHOMO_TPROXY

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}"