#!/usr/bin/env bash

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

function cmd::block::on_load() {
  flag::register --name
  flag::register --identity
  flag::register --type
  flag::register --force
  flag::register --quiet
  flag::register --ip
  flag::register --port
  flag::register --proto
  flag::register --subnet
  flag::register --block-name
  flag::register --service
}

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

function cmd::block::help() {
  cat <<EOF
Usage: wgctl block --name <name> [options]
   or: wgctl block --identity <identity> [options]

Block a client entirely or restrict access to specific IPs/ports/subnets/services.
All block rules are persisted and restored on WireGuard restart.
Clients blocked via --ip/--port/--subnet/--service remain in WireGuard
but have specific traffic restricted (shown as 'restricted' in list).

Options:
  --name <name>              Client name (e.g. phone-nuno)
  --identity <name>          Block all peers belonging to an identity
  --type <type>              Device type (optional, combines with --name)
  --ip <ip>                  Block access to specific IP (repeatable)
  --subnet <cidr>            Block access to subnet (repeatable)
  --port <ip:port:proto>     Block specific port (repeatable)
  --service <name>           Block a named service (repeatable)
  --block-name <name>        Optional name for this block rule
  --quiet                    Suppress output (used by group block)

Examples:
  wgctl block --name phone-nuno
  wgctl block --identity nuno
  wgctl block --name nuno --type phone
  wgctl block --name phone-nuno --ip 10.0.0.210
  wgctl block --name phone-nuno --service proxmox
  wgctl ban --name phone-nuno
EOF
}

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

function cmd::block::run() {
  local name="" identity="" type="" block_name=""
  local ips=() subnets=() ports=() services=()
  local quiet=false force=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)       name="$2";           shift 2 ;;
      --identity)   identity="$2";       shift 2 ;;
      --type)       type="$2";           shift 2 ;;
      --ip)         ips+=("$2");         shift 2 ;;
      --block-name) block_name="$2";     shift 2 ;;
      --service)    services+=("$2");    shift 2 ;;
      --force)      force=true;          shift   ;;
      --quiet)      quiet=true;          shift   ;;
      --subnet)     subnets+=("$2");     shift 2 ;;
      --port)       ports+=("$2");       shift 2 ;;
      --help)       cmd::block::help;    return  ;;
      *)
        log::error "Unknown flag: $1"
        cmd::block::help
        return 1
        ;;
    esac
  done

  # --identity: block all peers for this identity
  if [[ -n "$identity" ]]; then
    cmd::block::_block_identity "$identity" "$quiet" \
      "${ips[@]+"${ips[@]}"}" || return 1
    return 0
  fi

  [[ -z "$name" ]] && {
    log::error "Missing required flag: --name or --identity"
    cmd::block::help
    return 1
  }

  name=$(peers::resolve_and_require "$name" "$type") || return 1

  local client_ip
  client_ip=$(peers::get_ip "$name") || return 1

  # Full block if no specific targets
  if [[ ${#ips[@]} -eq 0 && ${#ports[@]} -eq 0 && \
        ${#subnets[@]} -eq 0 && ${#services[@]} -eq 0 ]]; then
    if peers::is_blocked "$name"; then
      log::wg_warning "Client is already blocked: ${name}"
      return 0
    fi
    monitor::update_endpoint_cache
    cmd::block::_block_all "$name" "$client_ip" "$quiet"
    return 0
  fi

  # Specific rules — check if already fully blocked
  if block::has_file "$name"; then
    local direct
    direct=$(block::is_blocked_direct "$name")
    if [[ "$direct" == "true" ]]; then
      log::wg_warning "${name} is fully blocked — unblock first to add specific rules"
      return 1
    fi
  fi

  local changed=false

  # Block specific IPs
  for ip in "${ips[@]}"; do
    ip::require_valid "$ip"
    fw::block_ip "$client_ip" "$ip"
    block::add_rule "$name" "$client_ip" "ip" "${block_name:-}" "$ip"
    $quiet || log::wg_success "${ip} has been blocked for ${name}"
  done

  # Block specific subnets
  for subnet in "${subnets[@]}"; do
    ip::require_valid "$subnet"
    fw::block_subnet "$client_ip" "$subnet"
    block::add_rule "$name" "$client_ip" "subnet" "${block_name:-}" "$subnet"
    $quiet || log::wg_success "${subnet} has been blocked for ${name}"
  done

  # Block specific ports
  for entry in "${ports[@]}"; do
    local b_target b_port b_proto
    IFS=":" read -r b_target b_port b_proto <<< "$entry"
    ip::require_valid "$b_target"
    fw::block_port "$client_ip" "$b_target" "$b_port" "${b_proto:-tcp}"
    block::add_rule "$name" "$client_ip" "port" "${block_name:-}" \
      "$b_target" "$b_port" "${b_proto:-tcp}"
    $quiet || log::wg_success "${client_ip}:${b_port}:${b_proto:-tcp} has been blocked for ${name}"
  done

  # Block services
  for svc in "${services[@]}"; do
    local resolved_lines=()
    mapfile -t resolved_lines < <(net::resolve "$svc" 2>/dev/null)
    if [[ ${#resolved_lines[@]} -eq 0 ]]; then
      log::error "Service not found or has no ports: ${svc}"
      return 1
    fi

    local already_blocked=true
    for resolved in "${resolved_lines[@]}"; do
      if [[ "$resolved" == *:*:* ]]; then
        local b_ip b_port b_proto
        IFS=":" read -r b_ip b_port b_proto <<< "$resolved"
        fw::has_block_rule "$client_ip" "$b_ip" "$b_port" "$b_proto" 2>/dev/null || \
          { already_blocked=false; break; }
      else
        fw::has_block_rule "$client_ip" "$resolved" 2>/dev/null || \
          { already_blocked=false; break; }
      fi
    done

    if $already_blocked; then
      $quiet || log::wg_warning "${svc} is already blocked for ${name}"
      continue
    fi

    for resolved in "${resolved_lines[@]}"; do
      if [[ "$resolved" == *:*:* ]]; then
        local b_ip b_port b_proto
        IFS=":" read -r b_ip b_port b_proto <<< "$resolved"
        fw::block_port "$client_ip" "$b_ip" "$b_port" "$b_proto"
        block::add_rule "$name" "$client_ip" "port" "$svc" \
          "$b_ip" "$b_port" "$b_proto"
      else
        fw::block_ip "$client_ip" "$resolved"
        block::add_rule "$name" "$client_ip" "ip" "$svc" "$resolved"
      fi
    done

    changed=true
    $quiet || log::wg_success "${svc} has been blocked for ${name}"
  done

  [[ ${#ips[@]} -gt 0 || ${#ports[@]} -gt 0 || \
     ${#subnets[@]} -gt 0 ]] && changed=true

  if $changed; then
    local peer_rule
    peer_rule=$(peers::get_meta "$name" "rule")
    if [[ -n "$peer_rule" ]] && rule::exists "$peer_rule"; then
      fw::flush_peer "$client_ip"
      rule::apply "$peer_rule" "$client_ip" "$name"
      block::restore_rules_for "$name" "$client_ip"
    fi
  fi

  return 0
}

# ============================================
# Identity block
# ============================================

function cmd::block::_block_identity() {
  local identity_name="${1:-}" quiet="${2:-false}"
  shift 2 || true

  identity::require_exists "$identity_name" || return 1
  identity::require_has_peers "$identity_name" || return 1

  local peers blocked=0 failed=0
  peers=$(identity::peers "$identity_name")

  while IFS= read -r peer_name; do
    [[ -z "$peer_name" ]] && continue
    if peers::is_blocked "$peer_name"; then
      $quiet || log::wg_warning "${peer_name} is already blocked"
      continue
    fi
    local client_ip
    client_ip=$(peers::get_ip "$peer_name") || continue
    monitor::update_endpoint_cache
    if cmd::block::_block_all "$peer_name" "$client_ip" true; then
      blocked=$(( blocked + 1 ))
    else
      failed=$(( failed + 1 ))
    fi
  done <<< "$peers"

  log::ok "Blocked ${blocked} peer(s) for identity '${identity_name}'"
  [[ $failed -gt 0 ]] && log::warn "${failed} peer(s) failed to block"
  return 0
}

# ============================================
# Helpers
# ============================================

function cmd::block::_get_endpoint() {
  local name="$1" public_key="$2"
  local endpoint
  endpoint=$(monitor::endpoint_for_key "$public_key")
  if [[ -z "$endpoint" || "$endpoint" == "(none)" ]]; then
    endpoint=$(monitor::get_cached_endpoint "$name")
  fi
  echo "$endpoint"
}

function cmd::block::_block_all() {
  local name="${1:?name required}"
  local client_ip="${2:?client_ip required}"
  local quiet="${3:-false}"

  block::apply_full "$name" "$client_ip"
  block::set_direct "$name" "$client_ip" "true"

  $quiet || log::wg_success "${name} has been blocked."
}