#!/usr/bin/env bash

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

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

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

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

Block a client entirely or restrict access to specific IPs/ports/subnets.
Block rules are persisted and restored on WireGuard restart.

Options:
  --name <name>         Client name (e.g. phone-nuno)
  --ip <ip>             Block access to specific IP (repeatable)
  --subnet <cidr>       Block access to subnet (repeatable)
  --port <ip:port:proto> Block specific port, e.g. 10.0.0.210:9000:tcp (repeatable)

Examples:
  wgctl block --name phone-nuno
  wgctl block --name phone-nuno --ip 10.0.0.210
  wgctl block --name phone-nuno --subnet 10.0.0.0/24
  wgctl block --name phone-nuno --port 10.0.0.210:9000:tcp
  wgctl ban --name phone-nuno --ip 10.0.0.100 --ip 10.0.0.200
EOF
}

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

function cmd::block::get_client_ip() {
  local name="$1"
  local conf
  conf="$(ctx::clients)/${name}.conf"

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

  grep "^Address" "$conf" | awk '{print $3}' | cut -d'/' -f1
}

# ============================================
# Block Run
# ============================================

function cmd::block::run() {
  local name=""
  local type=""
  local ips=()
  local subnets=()
  local ports=()

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)   name="$2";        shift 2 ;;
      --type)   type="$2";        shift 2 ;;
      --ip)     ips+=("$2");      shift 2 ;;
      --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

  if [[ -z "$name" ]]; then
    log::error "Missing required flag: --name"
    cmd::block::help
    return 1
  fi

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

  # Check if actually blocked
  if peers::is_blocked "$name" || [[ -f "$(ctx::block::path "${name}.block")" ]]; then
    log::wg_warning "Client is already blocked: ${name}"
    return 0
  fi

  # Update cache first
  monitor::update_endpoint_cache
  
  local public_key
  public_key=$(keys::public "$name") || return 1
  
  local endpoint
  endpoint=$(monitor::endpoint_for_key "$public_key")
  
  # Fall back to cache if live endpoint not available
  if [[ -z "$endpoint" || "$endpoint" == "(none)" ]]; then
    endpoint=$(monitor::get_cached_endpoint "$name")
  fi

  local client_ip
  client_ip=$(cmd::block::get_client_ip "$name") || return 1

  log::section "Blocking client: ${name} (${client_ip})"

  # No specific target — block everything
  if [[ ${#ips[@]} -eq 0 && ${#subnets[@]} -eq 0 && ${#ports[@]} -eq 0 ]]; then
      # Get real endpoint IP before removing peer from server
    local public_key endpoint
    public_key=$(keys::public "$name") || return 1
    endpoint=$(monitor::endpoint_for_key "$public_key")

    firewall::block_all "$client_ip" "$name"
    firewall::save_block "$name" "$client_ip"

    # Watch real endpoint IP if available
    if [[ -n "$endpoint" ]]; then
      monitor::unwatch "$client_ip"  # remove tunnel IP if added by block_all
      monitor::watch "$endpoint" "$name"
    fi

    # Remove peer from server to kill active connection
    peers::remove_from_server "$name"
    peers::reload

    log::wg_success "Client blocked: ${name}"
    return 0
  fi

  # Block specific IPs
  for ip in "${ips[@]}"; do
    ip::require_valid "$ip"
    firewall::block_ip "$client_ip" "$ip"
    firewall::save_block "$name" "$client_ip" "$ip"
  done

  # Block specific subnets
  for subnet in "${subnets[@]}"; do
    ip::require_valid "$subnet"
    firewall::block_subnet "$client_ip" "$subnet"
    firewall::save_block "$name" "$client_ip" "$subnet"
  done

  # Block specific ports
  for entry in "${ports[@]}"; do
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    ip::require_valid "$target"
    firewall::block_port "$client_ip" "$target" "$port" "$proto"
    firewall::save_block "$name" "$client_ip" "$target" "$port" "$proto"
  done

  log::wg_success "Block rules applied for: ${name}"
}