#!/usr/bin/env bash

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

function cmd::fw::on_load() {
  flag::register --peer
  flag::register --type
  flag::register --rule
  flag::register --no-nflog
  flag::register --no-accept
  flag::register --no-drop
}

# ============================================
# Help
# ============================================

function cmd::fw::help() {
  cat <<EOF
Usage: wgctl fw [subcommand] [options]

Inspect iptables firewall rules applied by wgctl.

Subcommands:
  list (default)              Show FORWARD chain rules
  nat                         Show NAT/PREROUTING rules (DNS redirects)
  flush-nat                   Flush NAT rules for a subnet
  count                       Show rule counts by type (ACCEPT/DROP/NFLOG)

Options for list:
  --peer <name>               Filter rules for a specific peer
  --rule <rule>               Filter by rule — shows all peers with that rule
  --no-nflog                  Hide NFLOG logging rules
  --no-accept                 Hide ACCEPT rules
  --no-drop                   Hide DROP rules

Examples:
  wgctl fw
  wgctl fw list --peer phone-nuno
  wgctl fw list --rule guest
  wgctl fw list --no-nflog
  wgctl fw list --peer phone-nuno --no-nflog
  wgctl fw nat
  wgctl fw count
  wgctl fw flush-nat --subnet 10.1.103.0/24
EOF
}

# ============================================
# Run
# ============================================

function cmd::fw::run() {
  local subcmd="${1:-list}"

  # If first arg is a flag, default to list
  if [[ "$subcmd" == --* ]]; then
    subcmd="list"
  else
    shift || true
  fi

  case "$subcmd" in
    list)      cmd::fw::list "$@"      ;;
    nat)       cmd::fw::nat "$@"       ;;
    flush-nat) cmd::fw::flush_nat "$@" ;;
    clean)     cmd::fw::clean          ;;
    count)     cmd::fw::count          ;;
    help)      cmd::fw::help           ;;
    *) log::error "Unknown subcommand: $subcmd"; return 1 ;;
  esac
}

function cmd::fw::list() {
  local peer="" type="" rule=""
  local show_nflog=true show_accept=true show_drop=true

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --peer)         peer="$2";          shift 2 ;;
      --rule)         rule="$2";          shift 2 ;;
      --type)         type="$2";          shift 2 ;;
      --no-nflog)     show_nflog=false;   shift   ;;
      --no-accept)    show_accept=false;  shift   ;;
      --no-drop)      show_drop=false;    shift   ;;
      *) shift ;;
    esac
  done

  # Rule filter — collect all IPs for peers with this rule
  if [[ -n "$rule" ]]; then
    log::section "Firewall Rules (FORWARD) — rule: ${rule}"
    printf "\n"
    local found=false
    $found || log::wg_warning "No peers found with rule: ${rule}"
    fw::list_peer_rules "$ip" "$show_nflog"
    printf "\n"
    return 0
  fi

  log::section "Firewall Rules (FORWARD)"
  printf "\n"

  if [[ -n "$peer" ]]; then
    local ip
    ip=$(peers::get_ip "$peer")
    [[ -z "$ip" ]] && log::error "Peer not found: $peer" && return 1
    iptables -L FORWARD -n -v | grep -F "$ip" \
      | cmd::fw::_print_filtered "$show_nflog" "$show_accept" "$show_drop" || true
  elif [[ -n "$type" ]]; then
    local subnet
    subnet=$(config::subnet_for "$type")
    [[ -z "$subnet" ]] && log::error "Unknown type: $type" && return 1
    iptables -L FORWARD -n -v | grep -F "$subnet" \
      | cmd::fw::_print_filtered "$show_nflog" "$show_accept" "$show_drop"
  else
    iptables -L FORWARD -n -v \
      | grep -v "^Chain\|^target\|^$\|ACCEPT.*0\.0\.0\.0.*0\.0\.0\.0" \
      | cmd::fw::_print_filtered "$show_nflog" "$show_accept" "$show_drop"
  fi

  printf "\n"
}

function cmd::fw::nat() {
  log::section "NAT Rules (PREROUTING)"
  printf "\n"
  iptables -t nat -L PREROUTING -n -v | while IFS= read -r rule; do
    [[ -z "$rule" ]] && continue
    ui::firewall_rule "$rule"
  done
  printf "\n"
}

function cmd::fw::flush_nat() {
  local type=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --type) type="$2"; shift 2 ;;
      *) shift ;;
    esac
  done

  if [[ -z "$type" ]]; then
    log::error "Missing required flag: --type"
    return 1
  fi

  local subnet
  subnet=$(config::subnet_for "$type")
  if [[ -z "$subnet" ]]; then
    log::error "Unknown type: $type"
    return 1
  fi

  local nat_linenums=()
  while IFS= read -r linenum; do
    [[ -n "$linenum" ]] && nat_linenums+=("$linenum")
  done < <(iptables -t nat -L PREROUTING -n --line-numbers | grep -F "${subnet}" | awk '{print $1}')

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

  log::wg_success "Flushed ${count} NAT rules for type '${type}' (${subnet}.0/24)"
}

function cmd::fw::stats() {
  log::section "Firewall Stats"
  iptables -L FORWARD -n -v | grep -v "^Chain\|^target\|^$" | \
    awk 'NR>1 {printf "  %8s pkts  %8s bytes  %s\n", $1, $2, $0}'
}

function cmd::fw::count() {
  log::section "Firewall Rule Counts"
  local drop accept nflog total
  drop=$(iptables -L FORWARD -n | grep -c "^DROP" || echo 0)
  accept=$(iptables -L FORWARD -n | grep -c "^ACCEPT" || echo 0)
  nflog=$(iptables -L FORWARD -n | grep -c "^NFLOG" || echo 0)
  total=$(( drop + accept + nflog ))
  printf "\n  %-10s %s\n" "DROP:"   "$drop"
  printf "  %-10s %s\n" "ACCEPT:" "$accept"
  printf "  %-10s %s\n" "NFLOG:"  "$nflog"
  printf "  %-10s %s\n" "TOTAL:"  "$total"
  printf "\n"
}

function cmd::fw::clean() {
  log::section "Duplicate Rule Report"
  iptables-save | sort | uniq -d | grep "^-A FORWARD"
}

function cmd::fw::_print_filtered() {
  local show_nflog="$1" show_accept="$2" show_drop="$3"
  while IFS= read -r rule; do
    [[ -z "$rule" ]] && continue
    if ! $show_nflog  && [[ "$rule" =~ NFLOG  ]]; then continue; fi
    if ! $show_accept && [[ "$rule" =~ ACCEPT ]]; then continue; fi
    if ! $show_drop   && [[ "$rule" =~ DROP   ]]; then continue; fi
    ui::firewall_rule "$rule"
  done
}