#!/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 WireGuard and firewall state matches expected configuration.

Options:
  --fix           Attempt to fix detected issues
  --peer <name>   Audit specific peer only
  --type <type>   Audit peers of specific type only

Examples:
  wgctl audit
  wgctl audit --peer phone-nuno
  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"
  local fix="$2"
  local actual="${3:-0}"
  local ip rule

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

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

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

  # Count expected iptables rules from rule file
  local rule_file
  rule_file="$(ctx::rule::path "${rule}.rule")"
  local block_ports block_ips allow_ports allow_ips expected
  block_ports=$(json::count "$rule_file" "block_ports")
  block_ips=$(json::count "$rule_file" "block_ips")
  allow_ports=$(json::count "$rule_file" "allow_ports")
  allow_ips=$(json::count "$rule_file" "allow_ips")
  expected=$(( (block_ports + block_ips) * 2 + allow_ports + allow_ips ))

  # actual is passed in as $3 — no iptables call here
  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"
      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"
      test::pass "  Fixed: $peer_name"
    fi
  fi
}

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
}