#!/usr/bin/env bash

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

function cmd::rule::on_load() {
  flag::register --name
  flag::register --desc
  flag::register --group
  flag::register --extends
  flag::register --remove-extends
  flag::register --block-ip
  flag::register --allow-ip
  flag::register --block-port
  flag::register --allow-port
  flag::register --remove-block-ip
  flag::register --remove-allow-ip
  flag::register --remove-block-port
  flag::register --remove-allow-port
  flag::register --peer
  flag::register --peers
  flag::register --dns-redirect
  flag::register --color
  flag::register --base
  flag::register --no-base
  flag::register --tree
  flag::register --resolved
  flag::register --force
  flag::register --type
}

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

function cmd::rule::help() {
  cat <<EOF
Usage: wgctl rule <subcommand> [options]

Manage firewall rules for peers.

Subcommands:
  list, ls                    List all rules
  show                        Show rule details
  inspect                     Show full inheritance tree
  add, new, create            Create a new rule
  update, edit                Update a rule and re-apply to all peers
  remove, rm, del             Remove a rule
  assign                      Assign a rule to a peer
  unassign                    Remove rule from a peer
  migrate                     Apply default rules to all unassigned peers
  reapply                     Re-apply a rule to all assigned peers

Options for list:
  --base                      Show base rules section
  --no-base                   Hide base rules section (default shows them)
  --group <name>              Filter by group name
  --tree                      Show full inheritance tree inline

Options for add/update:
  --name <name>               Rule name
  --desc <description>        Human readable description
  --group <group>             Display group (e.g. vm-rules, user-rules)
  --extends <rule,...>        Inherit from base rules (add only)
  --add-extends <rule,...>    Add base rules (update only)
  --remove-extends <rule,...> Remove base rules (update only)
  --allow-ip <ip/cidr>        Allow IP or subnet (repeatable)
  --allow-port <ip:port:proto> Allow specific port (repeatable)
  --block-ip <ip/cidr>        Block IP or subnet (repeatable)
  --block-port <ip:port:proto> Block specific port (repeatable)
  --dns-redirect              Force DNS through Pi-hole
  --remove-allow-ip <ip>      Remove allow IP entry
  --remove-allow-port <entry> Remove allow port entry
  --remove-block-ip <ip>      Remove block IP entry
  --remove-block-port <entry> Remove block port entry

Options for inspect:
  --name <name>               Rule name
  --peers                     Show assigned peers
  --resolved                  Show resolved/merged rule entries

Options for assign/unassign:
  --name <rule>               Rule name
  --peer <peer>               Peer name
  --type <type>               Peer device type (optional)

Examples:
  wgctl rule list
  wgctl rule list --base
  wgctl rule list --group vm-rules
  wgctl rule list --tree
  wgctl rule show --name guest
  wgctl rule inspect --name moonlight-02
  wgctl rule inspect --name moonlight-02 --resolved --peers
  wgctl rule add --name dev-01 --desc "Dev VM" --group vm-rules --extends no-lan
  wgctl rule update --name dev-01 --add-extends no-nginx
  wgctl rule update --name dev-01 --remove-extends no-nginx
  wgctl rule update --name dev-01 --group infra-rules
  wgctl rule assign --name dev-01 --peer laptop-nuno
  wgctl rule unassign --peer laptop-nuno
EOF
}

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

function cmd::rule::run() {
  local subcmd="${1:-help}"
  shift || true

  case "$subcmd" in
    list|ls)              cmd::rule::list "$@"     ;;
    show|inspect)         cmd::rule::show "$@"     ;;
    inspect)              cmd::rule::inspect "$@"  ;;
    add|new|create)       cmd::rule::add "$@"      ;;
    update|edit)          cmd::rule::update "$@"   ;;
    remove|rm|del|delete) cmd::rule::remove "$@"   ;;
    assign)               cmd::rule::assign "$@"   ;;
    unassign)             cmd::rule::unassign "$@" ;;
    migrate)              cmd::rule::migrate "$@"  ;;
    reapply)              cmd::rule::reapply "$@"  ;;
    help)                 cmd::rule::help          ;;
    *)
      log::error "Unknown subcommand: '${subcmd}'"
      cmd::rule::help
      return 1
      ;;
  esac
}

# ============================================
# List
# ============================================

function cmd::rule::_pad() {
  local text="$1" width="$2"
  local visible
  visible=$(printf "%s" "$text" | sed 's/\x1b\[[0-9;]*m//g')
  local visible_len=${#visible}
  local byte_len=${#text}
  local extra=$(( byte_len - visible_len ))
  printf "%-$(( width + extra ))s" "$text"
}

function cmd::rule::_print_extends_tree() {
  local extends="$1" indent="${2:-2}" rules_dir="$3"
  [[ -z "$extends" ]] && return 0

  local extend_list=()
  IFS=',' read -ra extend_list <<< "$extends"

  for base in "${extend_list[@]}"; do
    [[ -z "$base" ]] && continue
    local spaces
    spaces=$(printf '%*s' "$indent" '')
    printf "  \033[0;37m%s↳ %s\033[0m\n" "$spaces" "$base"

    if [[ "$indent" -lt 12 ]]; then
      local sub_file=""
      if sub_file=$(json::find_rule_file "$rules_dir" "$base" 2>/dev/null); then
        local sub_extends=""
        sub_extends=$(json::get "$sub_file" "extends" 2>/dev/null \
          | tr '\n' ',' | sed 's/,$//' || true)
        if [[ -n "$sub_extends" ]]; then
          cmd::rule::_print_extends_tree \
            "$sub_extends" $(( indent + 4 )) "$rules_dir"
        fi
      fi
    fi
  done
}

function cmd::rule::list() {
  local rules_dir
  rules_dir="$(ctx::rules)"

  local show_base_only=false
  local show_base=true
  local filter_group=""
  local show_tree=false
  local found_any=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --base)    show_base_only=true;  shift   ;;
      --no-base) show_base=false;      shift   ;;
      --group)   util::require_flag "--group" "${2:-}" || return 1
                 filter_group="${2,,}"; shift 2 ;;
      --tree)    show_tree=true;       shift   ;;
      --help)    cmd::rule::help;      return  ;;
      *)
        log::error "Unknown flag: $1"
        return 1 ;;
    esac
  done

  local rules=("${rules_dir}"/*.rule)
  if [[ ! -f "${rules[0]}" ]]; then
    log::wg "No rules configured"
    return 0
  fi

  local header_printed=false
  local printing_base=false
  local current_group=""

  while IFS="|" read -r name desc n_allows n_blocks \
        peer_count extends is_base group; do
    [[ -z "$name" ]] && continue

    # --base: show ONLY base rules
    if $show_base_only && [[ "$is_base" == "False" ]]; then
      continue
    fi

    # --no-base: hide base rules
    if ! $show_base && [[ "$is_base" == "True" ]]; then
      continue
    fi

    # --group filter (case insensitive)
    if [[ -n "$filter_group" && "${group,,}" != "$filter_group" ]]; then
      continue
    fi

    # Print header on first match
    if ! $header_printed; then
      log::section "Firewall Rules"
      printf "\n  %-20s %-40s %-8s %-8s %s\n" \
        "NAME" "DESCRIPTION" \
        "$(ui::center "ALLOWS" 8)" \
        "$(ui::center "BLOCKS" 8)" \
        "PEERS"
      local divider
      divider=$(printf '─%.0s' {1..88})
      printf "  %s\n" "$divider"
      header_printed=true
    fi

    # Base rules section header
    if [[ "$is_base" == "True" && "$printing_base" == "false" ]]; then
      if ! $show_base_only; then
        local bdashes
        bdashes=$(printf '─%.0s' {1..74})
        printf "\n  \033[0;37m── Base Rules \033[0m%s\n" "$bdashes"
      fi
      printing_base=true
      current_group=""
    fi

    # Group header — only for non-base rules
    if [[ "$is_base" == "False" && "${group,,}" != "${current_group,,}" ]]; then
      if [[ -n "$group" ]]; then
        printf "\n  \033[0;36m▸ %s\033[0m\n" "$group"
      elif [[ -n "$current_group" ]]; then
        printf "\n"
      fi
      current_group="$group"
    fi

    local short_desc="${desc:0:35}"
    [[ ${#desc} -gt 35 ]] && short_desc="${short_desc}..."

    local desc_col_width=40
    [[ "${short_desc:-—}" == "—" ]] && desc_col_width=42

    found_any=true

    printf "  %-20s %-${desc_col_width}s %-8s %-8s %s\n" \
      "$name" "${short_desc:-—}" \
      "$(ui::center "$n_allows" 8)" \
      "$(ui::center "$n_blocks" 8)" \
      "${peer_count} peers"

    # Print extends
    if [[ -n "$extends" ]]; then
      if $show_tree; then
        cmd::rule::_print_extends_tree "$extends" 2 "$rules_dir"
      else
        local extend_list=()
        IFS=',' read -ra extend_list <<< "$extends"
        for base in "${extend_list[@]}"; do
          [[ -z "$base" ]] && continue
          printf "  \033[0;37m  ↳ %s\033[0m\n" "$base"
        done
      fi
      printf "\n"
    fi

  done < <(json::rule_list_data "$rules_dir" "$(ctx::meta)")

  if ! $found_any; then
    if [[ -n "$filter_group" ]]; then
      log::wg_warning "No rules found in group: ${filter_group}"
    else
      log::wg_warning "No rules found"
    fi
  fi

  printf "\n"
}

# ============================================
# Show
# ============================================

function cmd::rule::show() {
  local name="" show_peers=true color=false show_resolved=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)     util::require_flag "--name" "${2:-}" || return 1
                  name="$2"; shift 2 ;;
      --no-peers) show_peers=false; shift ;;
      --color)    color=true;       shift ;;
      --resolved) show_resolved=true; shift ;;
      --help)     cmd::rule::help;  return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" ]] && log::error "Missing required flag: --name" && return 1
  rule::require_exists "$name" || return 1

  local rule_file
  rule_file="$(rule::path "$name")"

  local dns_redirect
  dns_redirect=$(rule::get_own "$name" "dns_redirect")
  dns_redirect="${dns_redirect:-false}"

  local resolved_dns
  resolved_dns=$(rule::get "$name" "dns_redirect")
  resolved_dns="${resolved_dns:-false}"

  local dns_display
  if [[ "${resolved_dns,,}" == "true" && "${dns_redirect,,}" != "true" ]]; then
    dns_display="true (inherited)"
  elif [[ "${dns_redirect,,}" == "true" ]]; then
    dns_display="true"
  else
    dns_display="false"
  fi


  log::section "Rule: ${name}"
  printf "\n"

  # ── Header info ──────────────────────────────
  local desc group dns_redirect
  desc=$(json::get "$rule_file" "desc")
  group=$(json::get "$rule_file" "group")

  ui::row "Description" "${desc:-—}"
  ui::row "Group"       "${group:-—}"
  ui::row "DNS"         "$dns_display"

  # ── Extends section ──────────────────────────
  local extends_raw=()
  mapfile -t extends_raw < <(json::get "$rule_file" "extends" 2>/dev/null || true)

  if [[ ${#extends_raw[@]} -gt 0 && -n "${extends_raw[0]}" ]]; then
    cmd::rule::_show_section "Extends"

    for base_name in "${extends_raw[@]}"; do
      [[ -z "$base_name" ]] && continue
      printf "\n    \033[0;37m↳ %s\033[0m\n" "$base_name"

      # Show resolved entries for this base
      local base_allow_ports base_allow_ips base_block_ips base_block_ports base_dns
      base_allow_ports=$(rule::get "$base_name" "allow_ports" 2>/dev/null || true)
      base_allow_ips=$(rule::get "$base_name" "allow_ips"    2>/dev/null || true)
      base_block_ips=$(rule::get "$base_name" "block_ips"    2>/dev/null || true)
      base_block_ports=$(rule::get "$base_name" "block_ports" 2>/dev/null || true)
      base_dns=$(json::get "$(rule::path "$base_name")" "dns_redirect" 2>/dev/null || true)

      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;32m+\033[0m %s\n" "$e"
      done <<< "$base_allow_ports"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;32m+\033[0m %s\n" "$e"
      done <<< "$base_allow_ips"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;31m-\033[0m %s\n" "$e"
      done <<< "$base_block_ips"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;31m-\033[0m %s\n" "$e"
      done <<< "$base_block_ports"
      [[ "$base_dns" == "true" ]] && \
        printf "      \033[0;36m↺\033[0m DNS → %s\n" "$(config::dns)"
    done
  fi

  # ── Own Rules section ─────────────────────────
  local own_allow_ports own_allow_ips own_block_ips own_block_ports
  own_allow_ports=$(json::get "$rule_file" "allow_ports")
  own_allow_ips=$(json::get "$rule_file" "allow_ips")
  own_block_ips=$(json::get "$rule_file" "block_ips")
  own_block_ports=$(json::get "$rule_file" "block_ports")

  local has_own=false
  [[ -n "$own_allow_ports" || -n "$own_allow_ips" || \
     -n "$own_block_ips"   || -n "$own_block_ports" ]] && has_own=true

  if $has_own; then
    if [[ ${#extends_raw[@]} -gt 0 && -n "${extends_raw[0]}" ]]; then
      cmd::rule::_show_section "Own Rules"
      printf "\n"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;32m+\033[0m %s\n" "$e"
      done <<< "$own_allow_ports"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;32m+\033[0m %s\n" "$e"
      done <<< "$own_allow_ips"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;31m-\033[0m %s\n" "$e"
      done <<< "$own_block_ips"
      while IFS= read -r e; do
        [[ -z "$e" ]] && continue
        printf "      \033[0;31m-\033[0m %s\n" "$e"
      done <<< "$own_block_ports"
      [[ "$dns_redirect" == "true" ]] && \
        printf "      \033[0;36m↺\033[0m DNS → %s\n" "$(config::dns)"
    else
      # No inheritance — use Allow/Block sections
       if [[ -n "$own_allow_ports" || -n "$own_allow_ips" ]]; then
        cmd::rule::_show_section "Allow"
        printf "\n"
        while IFS= read -r e; do
          [[ -z "$e" ]] && continue
          printf "    \033[0;32m+\033[0m %s\n" "$e"
        done <<< "$own_allow_ports"
        while IFS= read -r e; do
          [[ -z "$e" ]] && continue
          printf "    \033[0;32m+\033[0m %s\n" "$e"
        done <<< "$own_allow_ips"
      fi

      if [[ -n "$own_block_ips" || -n "$own_block_ports" ]]; then
        cmd::rule::_show_section "Block"
        printf "\n"
        while IFS= read -r e; do
          [[ -z "$e" ]] && continue
          printf "    \033[0;31m-\033[0m %s\n" "$e"
        done <<< "$own_block_ips"
        while IFS= read -r e; do
          [[ -z "$e" ]] && continue
          printf "    \033[0;31m-\033[0m %s\n" "$e"
        done <<< "$own_block_ports"
      fi

      if [[ "$dns_redirect" == "true" ]]; then
        cmd::rule::_show_section "DNS"
        printf "\n    \033[0;36m↺\033[0m Redirect all DNS → %s\n" "$(config::dns)"
      fi
    fi
  elif [[ ${#extends_raw[@]} -eq 0 || -z "${extends_raw[0]}" ]]; then
    printf "\n"
    ui::row "Access" "full (no restrictions)"
  fi

  # ── Resolved section (optional) ──────────────
  if $show_resolved; then
    cmd::rule::_show_section "Resolved (applied to peers)"
    local res_allow_ports res_allow_ips res_block_ips res_block_ports
    res_allow_ports=$(rule::get "$name" "allow_ports")
    res_allow_ips=$(rule::get "$name" "allow_ips")
    res_block_ips=$(rule::get "$name" "block_ips")
    res_block_ports=$(rule::get "$name" "block_ports")

    while IFS= read -r e; do [[ -n "$e" ]] && \
      printf "    \033[0;32m+\033[0m %s\n" "$e"; done <<< "$res_allow_ports"
    while IFS= read -r e; do [[ -n "$e" ]] && \
      printf "    \033[0;32m+\033[0m %s\n" "$e"; done <<< "$res_allow_ips"
    while IFS= read -r e; do [[ -n "$e" ]] && \
      printf "    \033[0;31m-\033[0m %s\n" "$e"; done <<< "$res_block_ips"
    while IFS= read -r e; do [[ -n "$e" ]] && \
      printf "    \033[0;31m-\033[0m %s\n" "$e"; done <<< "$res_block_ports"
  fi

  # ── Peers section ─────────────────────────────
  cmd::rule::_show_section "Peers"

  local peer_list=()
  mapfile -t peer_list < <(peers::with_rule "$name")
  local peer_count=${#peer_list[@]}
  ui::row "Assigned" "$peer_count"

  if $show_peers && [[ $peer_count -gt 0 ]]; then
    printf "\n"
    for peer_name in "${peer_list[@]}"; do
      local ip
      ip=$(peers::get_ip "$peer_name")
      printf "    %-28s %s\n" "$peer_name" "$ip"
    done
  fi

  printf "\n"
}

# ============================================
# Inspect
# ============================================

function cmd::rule::inspect() {
  local name="" show_peers=false show_resolved=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)     util::require_flag "--name" "${2:-}" || return 1
                  name="$2"; shift 2 ;;
      --peers)    show_peers=true;    shift ;;
      --resolved) show_resolved=true; shift ;;
      --help)     cmd::rule::help;    return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" ]] && log::error "Missing required flag: --name" && return 1
  rule::require_exists "$name" || return 1

  log::section "Rule Inspect: ${name}"

  local prev_section=""
  local show_resolved_flag="$show_resolved"

  while IFS="|" read -r section key value; do
    [[ -z "$section" ]] && continue

    # Skip resolved section unless requested
    if [[ "$section" == "resolved" ]] && ! $show_resolved_flag; then
      continue
    fi

    if [[ "$section" != "$prev_section" ]]; then
      case "$section" in
        own)         cmd::rule::_show_section "Own Rules"               ;;
        dns)         cmd::rule::_show_section "DNS"                     ;;
        resolved)    cmd::rule::_show_section "Resolved (applied)"      ;;
        inherited:*)
          local base_name="${section#inherited:}"
          cmd::rule::_show_section "Inherited: ${base_name}"            ;;
      esac
      prev_section="$section"
    fi

    case "$key" in
      allow_ip|allow_port)
        printf "    \033[0;32m+\033[0m %s\n" "$value" ;;
      block_ip|block_port)
        printf "    \033[0;31m-\033[0m %s\n" "$value" ;;
      dns_redirect)
        printf "    Redirect all DNS → %s\n" "$(config::dns)" ;;
    esac

  done < <(json::rule_inspect "$(ctx::rules)" "$name")

  if $show_peers; then
    cmd::rule::_show_section "Peers" "white" false
    local peer_list=()
    mapfile -t peer_list < <(peers::with_rule "$name")
    ui::row "Assigned" "${#peer_list[@]}"
    if [[ ${#peer_list[@]} -gt 0 ]]; then
      printf "\n"
      for peer_name in "${peer_list[@]}"; do
        local ip
        ip=$(peers::get_ip "$peer_name")
        printf "    %-28s %s\n" "$peer_name" "$ip"
      done
    fi
  fi

  printf "\n"
}

# ============================================
# Add
# ============================================

function cmd::rule::add() {
  local name="" desc="" group=""
  local extends=()
  local allow_ips=() block_ips=() block_ports=() allow_ports=()
  local dns_redirect=false
  local is_base=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)         name="$2";              shift 2 ;;
      --desc)         desc="$2";              shift 2 ;;
      --group)        group="$2";             shift 2 ;;
      --extends)
        IFS=',' read -ra ext <<< "$2"
        extends+=("${ext[@]}")
        shift 2 ;;
      --base)         is_base=true;           shift   ;;
      --allow-ip)     allow_ips+=("$2");      shift 2 ;;
      --allow-port)   allow_ports+=("$2");    shift 2 ;;
      --block-ip)     block_ips+=("$2");      shift 2 ;;
      --block-port)   block_ports+=("$2");    shift 2 ;;
      --dns-redirect) dns_redirect=true;      shift   ;;
      --help)         cmd::rule::help;        return  ;;
      *)
        log::error "Unknown flag: $1"
        return 1 ;;
    esac
  done

  [[ -z "$name" ]] && log::error "Missing required flag: --name" && return 1

  if rule::exists "$name"; then
    log::error "Rule already exists: ${name}"
    return 1
  fi

  # Validate extends
  for ext in "${extends[@]}"; do
    rule::require_exists "$ext" || return 1
  done

  # Determine target directory
  local rule_dir
  if $is_base; then
    rule_dir="$(ctx::rules)/base"
    mkdir -p "$rule_dir"
  else
    rule_dir="$(ctx::rules)"
  fi

  local rule_file="${rule_dir}/${name}.rule"

  local allow_str block_str port_str allow_port_str extends_str
  allow_str=$(IFS=','; echo "${allow_ips[*]}")
  block_str=$(IFS=','; echo "${block_ips[*]}")
  port_str=$(IFS=','; echo "${block_ports[*]}")
  allow_port_str=$(IFS=','; echo "${allow_ports[*]}")
  extends_str=$(IFS=','; echo "${extends[*]}")

  json::create_rule "$rule_file" "$name" "$desc" \
    "$($dns_redirect && echo true || echo false)" \
    "$allow_str" "$block_str" "$port_str" \
    "$allow_port_str" "$extends_str" "$group" || return 1

  local base_label=""
  $is_base && base_label=" (base)"

  log::wg_success "Rule created: ${name}${base_label}"
}

# ============================================
# Update
# ============================================

function cmd::rule::update() {
  local name="" desc="" group=""
  local add_extends=() rm_extends=()
  local allow_ips=() block_ips=() block_ports=() allow_ports=()
  local rm_allow_ips=() rm_block_ips=() rm_block_ports=() rm_allow_ports=()
  local dns_redirect=""

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)              name="$2";                shift 2 ;;
      --desc)              desc="$2";                shift 2 ;;
      --group)             group="$2";               shift 2 ;;
      --add-extends)
        IFS=',' read -ra ext <<< "$2"
        add_extends+=("${ext[@]}")
        shift 2 ;;
      --remove-extends)
        IFS=',' read -ra ext <<< "$2"
        rm_extends+=("${ext[@]}")
        shift 2 ;;
      --allow-ip)          allow_ips+=("$2");        shift 2 ;;
      --allow-port)        allow_ports+=("$2");      shift 2 ;;
      --block-ip)          block_ips+=("$2");        shift 2 ;;
      --block-port)        block_ports+=("$2");      shift 2 ;;
      --remove-allow-ip)   rm_allow_ips+=("$2");     shift 2 ;;
      --remove-block-ip)   rm_block_ips+=("$2");     shift 2 ;;
      --remove-block-port) rm_block_ports+=("$2");   shift 2 ;;
      --remove-allow-port) rm_allow_ports+=("$2");   shift 2 ;;
      --dns-redirect)      dns_redirect=true;        shift   ;;
      --help)              cmd::rule::help;          return  ;;
      *)
        log::error "Unknown flag: $1"
        return 1 ;;
    esac
  done

  [[ -z "$name" ]] && log::error "Missing required flag: --name" && return 1
  rule::require_exists "$name" || return 1

  local rule_file
  rule_file="$(rule::path "$name")"

  # Update simple fields
  [[ -n "$desc"         ]] && json::set "$rule_file" "desc"         "\"$desc\""
  [[ -n "$group"        ]] && json::set "$rule_file" "group"        "\"$group\""
  [[ -n "$dns_redirect" ]] && json::set "$rule_file" "dns_redirect" "true"

  # Add entries
  for ip in "${allow_ips[@]}";   do json::append "$rule_file" "allow_ips"   "$ip"; done
  for ip in "${block_ips[@]}";   do json::append "$rule_file" "block_ips"   "$ip"; done
  for p  in "${block_ports[@]}"; do json::append "$rule_file" "block_ports" "$p";  done
  for p  in "${allow_ports[@]}"; do json::append "$rule_file" "allow_ports" "$p";  done

  # Add/remove extends
  for ext in "${add_extends[@]}"; do
    rule::require_exists "$ext" || return 1
    json::append "$rule_file" "extends" "$ext"
  done
  for ext in "${rm_extends[@]}"; do
    json::remove "$rule_file" "extends" "$ext"
  done

  # Remove entries
  for ip in "${rm_allow_ips[@]}";   do json::remove "$rule_file" "allow_ips"   "$ip"; done
  for ip in "${rm_block_ips[@]}";   do json::remove "$rule_file" "block_ips"   "$ip"; done
  for p  in "${rm_block_ports[@]}"; do json::remove "$rule_file" "block_ports" "$p";  done
  for p  in "${rm_allow_ports[@]}"; do json::remove "$rule_file" "allow_ports" "$p";  done

  log::wg_success "Rule updated: ${name}"
  rule::reapply_all "$name"
}

# ============================================
# Remove / Assign / Unassign / Migrate / Reapply
# (unchanged from original)
# ============================================

function cmd::rule::remove() {
  local name="" force=false

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)  name="$2";       shift 2 ;;
      --force) force=true;      shift   ;;
      --help)  cmd::rule::help; return  ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" ]] && log::error "Missing required flag: --name" && return 1
  rule::require_exists "$name" || return 1

  local peer_list=()
  mapfile -t peer_list < <(peers::with_rule "$name")
  local peer_count=${#peer_list[@]}

  if [[ "$peer_count" -gt 0 ]]; then
    log::error "Rule '${name}' is assigned to ${peer_count} peer(s) — unassign first or use --force"
    $force || return 1
    for peer in "${peer_list[@]}"; do
      local ip
      ip=$(peers::get_ip "$peer")
      rule::unapply "$name" "$ip"
    done
  fi

  rm -f "$(rule::path "$name")"
  log::wg_success "Rule removed: ${name}"
}

function cmd::rule::assign() {
  local name="" peer="" type=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --peer) peer="$2"; shift 2 ;;
      --type) type="$2"; shift 2 ;;
      --help) cmd::rule::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" || -z "$peer" ]] && \
    log::error "Missing required flags: --name and --peer" && return 1

  rule::require_exists "$name"     || return 1
  rule::require_assignable "$name" || return 1

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

  local existing_rule
  existing_rule=$(peers::get_meta "$peer" "rule")

  local ip
  ip=$(peers::get_ip "$peer")
  log::debug "rule::assign: peer=$peer ip=$ip"
  [[ -z "$ip" ]] && log::error "Could not resolve IP for: $peer" && return 1
  log::debug "assign: peer=$peer ip=$ip clients=$(ctx::clients)"

  if [[ -n "$existing_rule" && "$existing_rule" != "$name" ]]; then
    rule::unapply "$existing_rule" "$ip"
    log::wg "Removed existing rule '${existing_rule}' from: ${peer}"
  fi

  rule::apply "$name" "$ip"
  log::wg_success "Assigned rule '${name}' to: ${peer}"
}

function cmd::rule::unassign() {
  local peer="" type=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --peer) peer="$2"; shift 2 ;;
      --type) type="$2"; shift 2 ;;
      --help) cmd::rule::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$peer" ]] && log::error "Missing required flag: --peer" && return 1
  peer=$(peers::resolve_and_require "$peer" "$type") || return 1

  local existing_rule
  existing_rule=$(peers::get_meta "$peer" "rule")

  if [[ -z "$existing_rule" ]]; then
    log::wg_warning "Peer '${peer}' has no assigned rule"
    return 0
  fi

  local ip
  ip=$(peers::get_ip "$peer")
  rule::unapply "$existing_rule" "$ip"
  log::wg_success "Unassigned rule from: ${peer}"
}

function cmd::rule::migrate() {
  log::section "Migrating peers to default rules"
  local tmp
  tmp=$(mktemp)

  while IFS= read -r peer_name; do
    local existing
    existing=$(peers::get_meta "$peer_name" "rule")
    [[ -n "$existing" ]] && continue
    local default_rule
    default_rule=$(peers::default_rule "$peer_name")
    rule::exists "$default_rule" || continue
    local ip
    ip=$(peers::get_ip "$peer_name")
    echo "${peer_name} ${default_rule} ${ip}" >> "$tmp"
  done < <(peers::all)

  local count=0
  local lines=()
  mapfile -t lines < "$tmp"
  rm -f "$tmp"

  for line in "${lines[@]}"; do
    IFS=" " read -r peer_name default_rule ip <<< "$line"
    rule::apply "$default_rule" "$ip" "$peer_name" </dev/null
    (( count++ )) || true
  done

  log::wg_success "Migrated ${count} peers"
}

function cmd::rule::reapply() {
  local name=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
  [[ -z "$name" ]] && log::error "Missing --name" && return 1
  rule::require_exists "$name" || return 1
  rule::reapply_all "$name"
  log::wg_success "Rule '${name}' reapplied"
}

# ============================================
# Show helpers
# ============================================

function cmd::rule::_show_section() {
  local title="${1:-}" color="${2:-white}" use_color="${3:-false}"
  local color_code=""
  if $use_color; then
    case "$color" in
      green) color_code="\033[0;32m" ;;
      red)   color_code="\033[0;31m" ;;
    esac
  fi
  printf "\n  ${color_code}── %s ──────────────────────────────────\033[0m\n" "$title"
}