#!/usr/bin/env bash

function cmd::audit::on_load() {
  flag::register --fix
  flag::register --peer
  flag::register --type
}

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

Verify that all peers have the correct iptables firewall rules applied.
Checks expected rule count (including inherited rules and peer-specific blocks)
vs actual iptables state.

Options:
  --peer <name>     Audit specific peer only
  --type <type>     Audit peers of specific device type only
  --fix             Attempt to auto-repair missing or extra rules

Output:
  ✅ pass   — peer has correct rule count
  ❌ fail   — peer has missing rules (run --fix to repair)
  ⚠️  warn   — peer has extra rules or configuration issues

Notes:
  Fully blocked peers (removed from WireGuard) show 0 expected fw rules.
  Restricted peers (peer-specific blocks) have their block rules included
  in the expected count.

Examples:
  wgctl audit
  wgctl audit --peer phone-nuno
  wgctl audit --type guest
  wgctl audit --fix
EOF
}

function cmd::audit::run() {
  local fix=false peer="" type=""

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --fix)   fix=true;   shift   ;;
      --peer)
        if [[ -z "${2:-}" ]]; then
          log::error "Flag --peer requires a value"
          cmd::audit::help
          return 1
        fi
        peer="$2"
        shift 2
        ;;
      --type)  type="$2";  shift 2 ;;
      --help)  cmd::audit::help; return ;;
      *)
        log::error "Unknown flag: $1"
        return 1
        ;;
    esac
  done

  test::reset
  log::section "WireGuard Audit"

  # Precompute iptables counts via single Python call
  declare -A peer_fw_counts
  while IFS=":" read -r peer_name count; do
    [[ -n "$peer_name" ]] && peer_fw_counts["$peer_name"]="$count"
  done < <(json::audit_fw_counts "$(ctx::clients)")

  test::section "Peer Rules"
  while IFS= read -r peer_name; do
    [[ -n "$peer" && "$peer_name" != "$peer" ]] && continue
    [[ -n "$type" && "$peer_name" != "${type}-"* ]] && continue
    cmd::audit::check_peer "$peer_name" "$fix" "${peer_fw_counts[$peer_name]:-0}"
  done < <(peers::all)

  test::section "WireGuard Server"
  cmd::audit::check_wg "$fix" "$peer" "$type"

  test::section "Rules"
  cmd::audit::check_rules

  test::summary
}

function cmd::audit::check_peer() {
  local peer_name="$1" fix="$2" actual="${3:-0}"
  local ip rule

  ip=$(peers::get_ip "$peer_name")
  rule=$(peers::get_meta "$peer_name" "rule")

  if [[ -z "$rule" ]]; then
    test::warn "$peer_name — no rule assigned (effective: $(peers::default_rule "$peer_name"))"
    return 0
  fi

  if ! rule::exists "$rule"; then
    test::fail "$peer_name — assigned rule '$rule' does not exist"
    return 0
  fi

  # Blocked peers have no fw rules (removed from wg0)
  if block::is_blocked "$peer_name" 2>/dev/null; then
    if [[ "$actual" -eq 0 ]]; then
      test::pass "$(printf "%-28s rule=%-15s blocked (no fw rules)" \
        "$peer_name" "$rule")"
    else
      test::warn "$(printf "%-28s rule=%-15s blocked but has %s fw rules" \
        "$peer_name" "$rule" "$actual")"
    fi
    return 0
  fi

  # Count expected rules from assigned rule (includes inheritance)
  local block_ports block_ips allow_ports allow_ips
  block_ports=$(json::count_resolved "$rule" "block_ports")
  block_ips=$(json::count_resolved "$rule" "block_ips")
  allow_ports=$(json::count_resolved "$rule" "allow_ports")
  allow_ips=$(json::count_resolved "$rule" "allow_ips")
  local expected=$(( (block_ports + block_ips) * 2 + allow_ports + allow_ips ))

  # Add peer-specific block rules (each ip/port/subnet = 2 rules: NFLOG+DROP)
  if block::has_specific_rules "$peer_name" 2>/dev/null; then
    while IFS="|" read -r bname btype target port proto; do
      [[ -z "$btype" || "$btype" == "full" ]] && continue
      (( expected += 2 )) || true
    done < <(block::get_rules "$peer_name")
  fi

  if [[ "$actual" -eq "$expected" ]]; then
    test::pass "$(printf "%-28s rule=%-15s fw: %s/%s" \
      "$peer_name" "$rule" "$actual" "$expected")"
  elif [[ "$actual" -gt "$expected" ]]; then
    test::warn "$(printf "%-28s rule=%-15s fw: %s/%s (extra rules)" \
      "$peer_name" "$rule" "$actual" "$expected")"
    if $fix; then
      fw::flush_peer "$ip"
      rule::apply "$rule" "$ip" "$peer_name"
      block::restore_rules_for "$peer_name" "$ip"
      test::pass "  Fixed: $peer_name"
    fi
  else
    test::fail "$(printf "%-28s rule=%-15s fw: %s/%s (missing rules)" \
      "$peer_name" "$rule" "$actual" "$expected")"
    if $fix; then
      rule::apply "$rule" "$ip" "$peer_name"
      block::restore_rules_for "$peer_name" "$ip"
      test::pass "  Fixed: $peer_name"
    fi
  fi
  return 0
}

function cmd::audit::check_wg() {
  local fix="$1" filter_peer="$2" filter_type="$3"
  local wg_peers
  wg_peers=$(wg show wg0 peers 2>/dev/null)

  while IFS= read -r peer_name; do
    [[ -n "$filter_peer" && "$peer_name" != "$filter_peer" ]] && continue
    [[ -n "$filter_type" && "$peer_name" != "${filter_type}-"* ]] && continue

    local pub_key
    pub_key=$(keys::public "$peer_name" 2>/dev/null)
    [[ -z "$pub_key" ]] && test::fail "$peer_name — missing public key" && continue

    if peers::is_blocked "$peer_name"; then
      if echo "$wg_peers" | grep -q "$pub_key"; then
        test::fail "$peer_name — blocked but still in wg0"
        $fix && peers::remove_from_server "$peer_name" && peers::reload
      else
        test::pass "$peer_name — correctly blocked (not in wg0)"
      fi
    else
      if echo "$wg_peers" | grep -q "$pub_key"; then
        test::pass "$peer_name — present in wg0"
      else
        test::fail "$peer_name — missing from wg0"
        if $fix; then
          local ip
          ip=$(peers::get_ip "$peer_name")
          peers::add_to_server "$peer_name" "$pub_key" "$ip"
          peers::reload
        fi
      fi
    fi
  done < <(peers::all)
}

function cmd::audit::check_rules() {
  local rules_dir
  rules_dir="$(ctx::rules)"

  # Precompute peer counts
  declare -A rule_counts
  while IFS= read -r peer_name; do
    local r
    r=$(peers::get_meta "$peer_name" "rule")
    [[ -z "$r" ]] && continue
    rule_counts["$r"]=$(( ${rule_counts["$r"]:-0} + 1 ))
  done < <(peers::all)

  for rule_file in "${rules_dir}"/*.rule; do
    [[ -f "$rule_file" ]] || continue
    local name
    name=$(json::get "$rule_file" "name")
    local count=${rule_counts["$name"]:-0}
    test::pass "Rule '${name}' — ${count} peers assigned"
  done
}