#!/usr/bin/env bash

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

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

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

function firewall::block_ip() {
  local client_ip="$1"
  local target_ip="$2"

  iptables -A FORWARD -s "$client_ip" -d "$target_ip" -j LOG --log-prefix "wgctl-dropped: " --log-level 4
  iptables -A FORWARD -s "$client_ip" -d "$target_ip" -j DROP
  log::wg_block "Blocked ${client_ip} → ${target_ip}"
}

function firewall::unblock_ip() {
  local client_ip="$1"
  local target_ip="$2"

  iptables -D FORWARD -s "$client_ip" -d "$target_ip" -j LOG --log-prefix "wgctl-dropped: " --log-level 4 2>/dev/null || true
  iptables -D FORWARD -s "$client_ip" -d "$target_ip" -j DROP 2>/dev/null || true
  log::wg_unblock "Unblocked ${client_ip} → ${target_ip}"
}

function firewall::block_port() {
  local client_ip="$1"
  local target_ip="$2"
  local port="$3"
  local proto="${4:-tcp}"

  iptables -A FORWARD -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port" -j DROP
  log::wg_block "Blocked ${client_ip} → ${target_ip}:${port}/${proto}"
}

function firewall::unblock_port() {
  local client_ip="$1"
  local target_ip="$2"
  local port="$3"
  local proto="${4:-tcp}"

  iptables -D FORWARD -s "$client_ip" -d "$target_ip" -p "$proto" --dport "$port" -j DROP 2>/dev/null || true
  log::wg_unblock "Unblocked ${client_ip} → ${target_ip}:${port}/${proto}"
}

function firewall::block_all() {
  local client_ip="$1"
  local client_name="$2"

  iptables -A FORWARD -s "$client_ip" -j LOG --log-prefix "wgctl-dropped: " --log-level 4
  iptables -A FORWARD -s "$client_ip" -j DROP

  log::wg_block "Blocked all traffic from: ${client_ip}"
}

function firewall::unblock_all() {
  local client_ip="$1"

  iptables -D FORWARD -s "$client_ip" -j LOG --log-prefix "wgctl-dropped: " --log-level 4 2>/dev/null || true
  iptables -D FORWARD -s "$client_ip" -j DROP 2>/dev/null || true

  monitor::unwatch "$client_ip"
  log::wg_unblock "Unblocked all traffic from: ${client_ip}"
}

function firewall::block_subnet() {
  local client_ip="$1"
  local target_subnet="$2"

  iptables -A FORWARD -s "$client_ip" -d "$target_subnet" -j DROP
  log::wg_block "Blocked ${client_ip} → subnet ${target_subnet}"
}

function firewall::unblock_subnet() {
  local client_ip="$1"
  local target_subnet="$2"

  iptables -D FORWARD -s "$client_ip" -d "$target_subnet" -j DROP 2>/dev/null || true
  log::wg_unblock "Unblocked ${client_ip} → subnet ${target_subnet}"
}

# ============================================
# Guest Subnet Rules
# ============================================

# Sensitive services blocked for all guest peers
declare -ga GUEST_BLOCKED_SERVICES=(
  "10.0.0.100:8006:tcp"   # Proxmox UI
  "10.0.0.100:22:tcp"     # Proxmox SSH
  "10.0.0.105:8007:tcp"   # PBS UI
  "10.0.0.102:22:tcp"     # WireGuard LXC SSH
  "10.0.0.200:80:tcp"     # TrueNAS UI HTTP
  "10.0.0.200:443:tcp"    # TrueNAS UI HTTPS
  "10.0.0.103:80:tcp"     # Pi-hole WebUI
  "10.0.0.101:80:tcp"     # NPM WebUI HTTP
  "10.0.0.101:443:tcp"    # NPM WebUI HTTPS
  "10.0.0.210:9000:tcp"   # Portainer direct port
)

function firewall::guest_rules_applied() {
  local guest_subnet
  guest_subnet="$(config::subnet_for "guest").0/24"
  # Check if at least the first rule exists
  local first_entry="${GUEST_BLOCKED_SERVICES[0]}"
  local target port proto
  IFS=":" read -r target port proto <<< "$first_entry"
  proto="${proto:-tcp}"
  iptables -C FORWARD -s "$guest_subnet" -d "$target" -p "$proto" --dport "$port" -j DROP 2>/dev/null
}

function firewall::apply_guest_rules() {
  local guest_subnet
  guest_subnet="$(config::subnet_for "guest").0/24"

  # Skip if already applied
  if firewall::guest_rules_applied; then
    log::wg "Guest firewall rules already applied"
    return 0
  fi

  for entry in "${GUEST_BLOCKED_SERVICES[@]}"; do
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    iptables -I FORWARD 1 -s "$guest_subnet" -d "$target" -p "$proto" --dport "$port" -j DROP
    log::wg_block "Guest rule: blocked ${guest_subnet} → ${target}:${port}/${proto}"
  done

  # Persist guest rules marker
  local marker
  marker="$(ctx::blocks)/_guest_rules.active"
  touch "$marker"

  log::wg_block "Applied guest firewall rules"
  firewall::apply_guest_dns_redirect
}

function firewall::remove_guest_rules() {
  local guest_subnet
  guest_subnet="$(config::subnet_for "guest").0/24"

  for entry in "${GUEST_BLOCKED_SERVICES[@]}"; do
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    iptables -D FORWARD -s "$guest_subnet" -d "$target" -p "$proto" --dport "$port" -j DROP 2>/dev/null || true
  done

  # Remove persistence marker
  local marker
  marker="$(ctx::blocks)/_guest_rules.active"
  rm -f "$marker"

  log::wg_unblock "Removed guest firewall rules"
  firewall::remove_guest_dns_redirect
}

function firewall::apply_guest_dns_redirect() {
  if iptables -t nat -C PREROUTING -i wg0 -s "$(config::subnet_for guest).0/24" -p udp --dport 53 -j DNAT --to-destination "$(config::dns):53" 2>/dev/null; then
    log::wg "Guest DNS redirect already applied"
    return 0
  fi

  local guest_subnet dns
  guest_subnet="$(config::subnet_for "guest").0/24"
  dns="$(config::dns)"

  # Log DNS bypass attempts (queries not directed at Pi-hole)
  iptables -t nat -A PREROUTING -i wg0 -s "$guest_subnet" -p udp --dport 53 \
    ! -d "$dns" \
    -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4
  iptables -t nat -A PREROUTING -i wg0 -s "$guest_subnet" -p tcp --dport 53 \
    ! -d "$dns" \
    -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4

  # Redirect all DNS to Pi-hole
  iptables -t nat -A PREROUTING -i wg0 -s "$guest_subnet" -p udp --dport 53 -j DNAT --to-destination "${dns}:53"
  iptables -t nat -A PREROUTING -i wg0 -s "$guest_subnet" -p tcp --dport 53 -j DNAT --to-destination "${dns}:53"

  log::wg_block "Guest DNS redirected to Pi-hole (${dns}), bypass attempts will be logged"
}

function firewall::remove_guest_dns_redirect() {
  local guest_subnet dns
  guest_subnet="$(config::subnet_for "guest").0/24"
  dns="$(config::dns)"

  iptables -t nat -D PREROUTING -i wg0 -s "$guest_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 wg0 -s "$guest_subnet" -p tcp --dport 53 \
    ! -d "$dns" \
    -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4 2>/dev/null || true
  iptables -t nat -D PREROUTING -i wg0 -s "$guest_subnet" -p udp --dport 53 -j DNAT --to-destination "${dns}:53" 2>/dev/null || true
  iptables -t nat -D PREROUTING -i wg0 -s "$guest_subnet" -p tcp --dport 53 -j DNAT --to-destination "${dns}:53" 2>/dev/null || true

  log::wg_unblock "Removed guest DNS redirect"
}

# ============================================
# Persistence — block files
# ============================================

function firewall::save_block() {
  local name="$1"
  local client_ip="$2"
  local target="${3:-}"
  local port="${4:-}"
  local proto="${5:-}"

  local block_file
  block_file="$(ctx::block::path "${name}.block")"

  echo "${client_ip} ${target} ${port} ${proto}" >> "$block_file"
  log::wg_block "Persisted block rule for: ${name}"
}

function firewall::remove_block_file() {
  local name="$1"
  local block_file
  block_file="$(ctx::block::path "${name}.block")"

  rm -f "$block_file"
  log::wg_unblock "Removed block file for: ${name}"
}

function firewall::restore_blocks() {
  local blocks_dir
  blocks_dir="$(ctx::blocks)"

  # Restore guest rules if marker exists
  local marker="${blocks_dir}/_guest_rules.active"
  if [[ -f "$marker" ]]; then
    firewall::apply_guest_rules
    log::wg "Restored guest firewall rules"
  fi

  # Restore per-client block rules
  for block_file in "${blocks_dir}"/*.block; do
    [[ -f "$block_file" ]] || continue

    local name
    name=$(basename "$block_file" .block)

    while IFS=" " read -r client_ip target port proto; do
      if [[ -z "$target" ]]; then
        firewall::block_all "$client_ip" "$name"
      elif [[ -n "$port" ]]; then
        firewall::block_port "$client_ip" "$target" "$port" "${proto:-tcp}"
      else
        firewall::block_ip "$client_ip" "$target"
      fi
    done < "$block_file"

    log::wg "Restored block rules for: ${name}"
  done
}

# ============================================
# Preset Application
# ============================================

function firewall::apply_preset() {
  local name="$1"
  local client_ip="$2"
  local preset_file
  preset_file="$(ctx::preset::path "${name}.preset")"

  if [[ ! -f "$preset_file" ]]; then
    log::error "Preset not found: ${name}"
    return 1
  fi

  source "$preset_file"

  if [[ -n "${BLOCK_IPS:-}" ]]; then
    for ip in $BLOCK_IPS; do
      firewall::block_ip "$client_ip" "$ip"
      firewall::save_block "$client_ip" "$client_ip" "$ip"
    done
  fi

  if [[ -n "${BLOCK_SUBNETS:-}" ]]; then
    for subnet in $BLOCK_SUBNETS; do
      firewall::block_subnet "$client_ip" "$subnet"
      firewall::save_block "$client_ip" "$client_ip" "$subnet"
    done
  fi

  if [[ -n "${BLOCK_PORTS:-}" ]]; then
    for entry in $BLOCK_PORTS; do
      local target port proto
      IFS=":" read -r target port proto <<< "$entry"
      proto="${proto:-tcp}"
      firewall::block_port "$client_ip" "$target" "$port" "$proto"
      firewall::save_block "$name" "$client_ip" "$target" "$port" "$proto"
    done
  fi

  log::wg_preset "Applied preset '${name}' to: ${client_ip}"
}