#!/usr/bin/env bash

# ============================================
# Lifecycle
# ============================================

function fw::on_load() {
  system::require_command iptables
}

# ============================================
# Rule Management
# ============================================

# ============================================
# Public Interfaces

function fw::has_rule() {
  # Generic rule existence check
  # fw::has_rule <target> <client_ip> <dest> [port] [proto]
  local action="${1:-DROP}" client_ip="${2:-}" target="${3:-}" \
        port="${4:-}" proto="${5:-}"
  if [[ -n "$port" ]]; then
    fw::_forward_exists -s "$client_ip" -d "$target" \
      -p "$proto" --dport "$port" -j "$action"
  else
    fw::_forward_exists -s "$client_ip" -d "$target" -j "$action"
  fi
}

function fw::has_block_rule() {
  fw::has_rule "DROP"   "$@"
}

function fw::has_allow_rule() {
  fw::has_rule "ACCEPT" "$@"
}

# ============================================
# Block / Unblock

function fw::block_ip() {
  local client_ip="${1:-}" target_ip="${2:-}" mode="${3:-insert}"
  fw::_block_pair "$mode" -s "$client_ip" -d "$target_ip"
}

function fw::unblock_ip() {
  local client_ip="${1:-}" target_ip="${2:-}"
  fw::_unblock_pair -s "$client_ip" -d "$target_ip"
}

function fw::block_port() {
  local client_ip="${1:-}" target_ip="${2:-}" port="${3:-}" \
        proto="${4:-tcp}" mode="${5:-insert}"
  fw::_block_pair "$mode" \
    -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port"
}

function fw::unblock_port() {
  local client_ip="${1:-}" target_ip="${2:-}" port="${3:-}" proto="${4:-tcp}"
  fw::_unblock_pair \
    -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port"
}

function fw::block_subnet() {
  local client_ip="${1:-}" target_subnet="${2:-}" mode="${3:-append}"
  fw::_block_pair "$mode" -s "$client_ip" -d "$target_subnet"
  log::wg_block "Blocked ${client_ip} → subnet ${target_subnet}"
}

function fw::unblock_subnet() {
  local client_ip="${1:-}" target_subnet="${2:-}"
  fw::_unblock_pair -s "$client_ip" -d "$target_subnet"
  log::wg_unblock "Unblocked ${client_ip} → subnet ${target_subnet}"
}

function fw::block_all() {
  local client_ip="${1:-}" client_name="${2:-}"
  fw::_forward_exists -s "$client_ip" -j DROP \
    || iptables -A FORWARD -s "$client_ip" -j DROP
  log::debug "Blocked all traffic from: ${client_ip}"
}

function fw::unblock_all() {
  local client_ip="${1:-}"
  fw::_forward_exists -s "$client_ip" -j DROP \
    && iptables -D FORWARD -s "$client_ip" -j DROP 2>/dev/null || true
  monitor::unwatch "$client_ip"
  log::debug "Unblocked all traffic from: ${client_ip}"
}

# ============================================
# Allow / Unallow

function fw::allow_ip() {
  local client_ip="${1:-}" target_ip="${2:-}"
  fw::_accept_insert -s "$client_ip" -d "$target_ip"
}

function fw::unallow_ip() {
  local client_ip="${1:-}" target_ip="${2:-}"
  fw::_accept_remove -s "$client_ip" -d "$target_ip"
}

function fw::allow_subnet() {
  local client_ip="${1:-}" target_subnet="${2:-}"
  fw::_accept_insert -s "$client_ip" -d "$target_subnet"
}

function fw::unallow_subnet() {
  local client_ip="${1:-}" target_subnet="${2:-}"
  fw::_accept_remove -s "$client_ip" -d "$target_subnet"
}

function fw::allow_port() {
  local client_ip="${1:-}" target_ip="${2:-}" port="${3:-}" proto="${4:-tcp}"
  fw::_accept_insert \
    -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port"
}

function fw::unallow_port() {
  local client_ip="${1:-}" target_ip="${2:-}" port="${3:-}" proto="${4:-tcp}"
  fw::_accept_remove \
    -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port"
}

# ============================================
# Flush
# ============================================

function fw::flush_peer() {
  local client_ip="${1:?client_ip required}"
  log::debug "flush_peer: starting for $client_ip"

  local linenums=()
  while IFS= read -r linenum; do
    [[ -n "$linenum" ]] && linenums+=("$linenum")
  done < <(iptables -L FORWARD -n --line-numbers \
    | grep -F "$client_ip" | awk '{print $1}')

  local count=0
  local i
  for (( i=${#linenums[@]}-1; i>=0; i-- )); do
    iptables -D FORWARD "${linenums[$i]}" 2>/dev/null || true
    (( count++ )) || true
  done

  log::debug "flush_peer: removed $count FORWARD rules for: $client_ip"
}

function fw::flush_forward() {
  iptables -F FORWARD
  log::debug "Flushed FORWARD chain"
}

function fw::flush_nat() {
  iptables -t nat -F PREROUTING
  log::debug "Flushed NAT PREROUTING chain"
}

function fw::flush_all() {
  fw::flush_forward
  fw::flush_nat
}

# ============================================
# NAT / DNS Redirect
# ============================================

function fw::nat_add_dns_redirect() {
  local subnet="${1:-}" dns="${2:-}" interface="${3:-wg0}"
  iptables -t nat -A PREROUTING -i "$interface" -s "$subnet" \
    -p udp --dport 53 ! -d "$dns" \
    -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4
  iptables -t nat -A PREROUTING -i "$interface" -s "$subnet" \
    -p udp --dport 53 -j DNAT --to-destination "${dns}:53"
  iptables -t nat -A PREROUTING -i "$interface" -s "$subnet" \
    -p tcp --dport 53 -j DNAT --to-destination "${dns}:53"
}

function fw::nat_remove_dns_redirect() {
  local subnet="${1:-}" dns="${2:-}" interface="${3:-wg0}"
  iptables -t nat -D PREROUTING -i "$interface" -s "$subnet" \
    -p udp --dport 53 ! -d "$dns" \
    -j LOG --log-prefix "wgctl-dns-redirect: " \
    --log-level 4 2>/dev/null || true
  iptables -t nat -D PREROUTING -i "$interface" -s "$subnet" \
    -p udp --dport 53 -j DNAT \
    --to-destination "${dns}:53" 2>/dev/null || true
  iptables -t nat -D PREROUTING -i "$interface" -s "$subnet" \
    -p tcp --dport 53 -j DNAT \
    --to-destination "${dns}:53" 2>/dev/null || true
}

# ============================================
# Display
# ============================================

function fw::forward_rules_for_ip() {
  local ip="${1:-}"
  iptables -L FORWARD -n -v </dev/null | grep -F "$ip" || true
}

function fw::proto_name() {
  local proto="${1:-0}"
  case "$proto" in
    6)   echo "tcp"  ;;
    17)  echo "udp"  ;;
    1)   echo "icmp" ;;
    0)   echo "all"  ;;
    tcp|udp|icmp|all) echo "$proto" ;;
    *)   echo "$proto" ;;
  esac
}

function fw::format_rule() {
  local line="${1:-}"
  [[ -z "$line" ]] && return 0

  local target prot src dst extra
  target=$(awk '{print $3}' <<< "$line")
  prot=$(awk '{print $4}'   <<< "$line")
  src=$(awk '{print $8}'    <<< "$line")
  dst=$(awk '{print $9}'    <<< "$line")
  extra=$(awk '{for(i=10;i<=NF;i++) printf $i" "}' <<< "$line" | xargs)

  local prot_name
  prot_name=$(fw::proto_name "$prot")

  local dst_fmt="$dst"
  if [[ "$extra" =~ dpt:([0-9]+) ]]; then
    local port="${BASH_REMATCH[1]}"
    dst_fmt="${dst}:${port}:${prot_name}"
  fi

  local formatted
  formatted=$(printf "  %-8s %-15s → %s" "$target" "$src" "$dst_fmt")
  ui::firewall_rule "$formatted"
}

function fw::list_peer_rules() {
  local ip="${1:-}" show_nflog="${2:-false}"
  fw::forward_rules_for_ip "$ip" | while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    ! $show_nflog && [[ "$line" =~ NFLOG ]] && continue
    fw::format_rule "$line"
  done
}

# ============================================
# Counts
# ============================================

function fw::count_peer_rules() {
  local ip="${1:-}"
  local total=0 accepts=0 drops=0
  while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    [[ "$line" =~ NFLOG ]] && continue
    (( total++ )) || true
    [[ "$line" =~ ACCEPT ]] && (( accepts++ )) || true
    [[ "$line" =~ DROP   ]] && (( drops++   )) || true
  done < <(fw::forward_rules_for_ip "$ip")
  echo "${total}|${accepts}|${drops}"
}

# ============================================
# private
# ============================================

function fw::_forward_exists() {
  iptables -C FORWARD "$@" 2>/dev/null
}

function fw::_rule_exists() {
  local table="${1:-filter}" chain="${2:-FORWARD}"
  shift 2
  iptables -t "$table" -C "$chain" "$@" 2>/dev/null
}

function fw::_nat_exists() {
  fw::_rule_exists nat PREROUTING "$@"
}

# Core NFLOG+DROP block pair — insert or append
function fw::_block_pair() {
  local mode="${1:-insert}"  # insert | append
  shift
  # $@ = match args (no -j)
  if [[ "$mode" == "insert" ]]; then
    # insert: DROP first at pos 1, NFLOG second at pos 1 → NFLOG ends above DROP
    fw::_forward_exists "$@" -j DROP \
      || iptables -I FORWARD 1 "$@" -j DROP
    fw::_forward_exists "$@" -j NFLOG --nflog-group 1 --nflog-prefix "wgctl-drop:" \
      || iptables -I FORWARD 1 "$@" -j NFLOG --nflog-group 1 --nflog-prefix "wgctl-drop:"
  else
    # append: NFLOG first, DROP second → NFLOG ends above DROP
    fw::_forward_exists "$@" -j NFLOG --nflog-group 1 --nflog-prefix "wgctl-drop:" \
      || iptables -A FORWARD "$@" -j NFLOG --nflog-group 1 --nflog-prefix "wgctl-drop:"
    fw::_forward_exists "$@" -j DROP \
      || iptables -A FORWARD "$@" -j DROP
  fi
}

# Core NFLOG+DROP removal
function fw::_unblock_pair() {
  shift 0  # no mode needed for deletion
  # $@ = match args
  fw::_forward_exists "$@" -j NFLOG --nflog-group 1 --nflog-prefix "wgctl-drop:" \
    && iptables -D FORWARD "$@" -j NFLOG \
       --nflog-group 1 --nflog-prefix "wgctl-drop:" 2>/dev/null || true
  fw::_forward_exists "$@" -j DROP \
    && iptables -D FORWARD "$@" -j DROP 2>/dev/null || true
}

# Core ACCEPT insert
function fw::_accept_insert() {
  # $@ = match args
  fw::_forward_exists "$@" -j ACCEPT \
    || iptables -I FORWARD 1 "$@" -j ACCEPT
}

# Core ACCEPT removal
function fw::_accept_remove() {
  fw::_forward_exists "$@" -j ACCEPT \
    && iptables -D FORWARD "$@" -j ACCEPT 2>/dev/null || true
}