#!/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-service
  flag::register --allow-service
  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 --base
  flag::register --no-base
  flag::register --tree
  flag::register --detailed
  flag::register --resolved
  flag::register --force
  flag::register --type
  flag::register --all
}

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

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

Manage firewall rules with inheritance support.
Rules can extend base rules to compose reusable access policies.
Service names from 'wgctl net' can be used instead of raw IPs/ports.

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

Options for list:
  --base                      Show only base rules
  --no-base                   Hide base rules section
  --group <name>              Filter by group (case insensitive)
  --detailed                  Show rule entries inline

Options for add:
  --name <name>               Rule name
  --desc <description>        Description
  --group <group>             Display group (e.g. VM Rules, Users)
  --extends <rule,...>        Inherit from base rules (comma-separated)
  --base                      Create as base rule (not directly assignable)
  --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)
  --block-service <name>      Block named service (repeatable)
  --allow-service <name>      Allow named service (repeatable)
  --dns-redirect              Force DNS through Pi-hole

Options for update:
  (same as add, plus:)
  --add-extends <rule,...>    Add inherited base rules
  --remove-extends <rule,...> Remove inherited base rules
  --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 show:
  --name <name>               Rule name
  --resolved                  Show resolved/merged entries
  --no-peers                  Hide assigned peers

Options for reapply:
  --name <name>               Rule name
  --all                       Reapply all rules

Examples:
  wgctl rule list
  wgctl rule list --detailed
  wgctl rule list --group "VM Rules"
  wgctl rule show --name guest
  wgctl rule show --name moonlight-02 --resolved
  wgctl rule add --name no-proxmox --base --block-service proxmox
  wgctl rule add --name dev-01 --desc "Dev access" --extends no-lan
  wgctl rule assign --name dev-01 --peer laptop-nuno
  wgctl rule reapply --all
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    "$@" ;;
    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::list() {
  local rules_dir
  rules_dir="$(ctx::rules)"

  local show_base_only=false show_base=true
  local filter_group="" detailed=false

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

  local data
  data=$(json::rule_list_data "$rules_dir" "$(ctx::meta)")
  [[ -z "$data" ]] && log::wg "No rules configured" && return 0

  # Measure max name width
  local w_name=12
  while IFS='|' read -r name rest; do
    [[ -z "$name" ]] && continue
    (( ${#name} > w_name )) && w_name=${#name}
  done <<< "$data"
  (( w_name += 2 ))

  log::section "Firewall Rules"
  echo ""

  local current_group="" printing_base=false found_any=false

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

    $show_base_only && [[ "$is_base" == "False" ]] && continue
    ! $show_base    && [[ "$is_base" == "True"  ]] && continue
    [[ -n "$filter_group" && "${group,,}" != "$filter_group" ]] && continue

    found_any=true

    # Base rules section header
    if [[ "$is_base" == "True" && "$printing_base" == "false" ]]; then
      if ! $show_base_only; then
        ui::rule::list_base_header
      fi
      printing_base=true
      current_group=""
    fi

    # Group header — non-base rules only
    if [[ "$is_base" == "False" && "${group,,}" != "${current_group,,}" ]]; then
      if [[ -n "$group" ]]; then
        ui::rule::list_group_header "$group"
      elif [[ -n "$current_group" ]]; then
        echo ""
      fi
      current_group="$group"
    fi

    # Rule row
    # ui::rule::list_row "$name" "$n_allows" "$n_blocks" "$peer_count" "$w_name"

    # Extends
    # Rule row — pass extends_csv for compact inline display
    local compact_extends=""
    if [[ -z "$detailed" ]] || ! $detailed; then
      compact_extends="$extends"
    fi
    ui::rule::list_row "$name" "$n_allows" "$n_blocks" "$peer_count" "$w_name" "$compact_extends"
 
    # Detailed mode — show expanded entries
    if $detailed && [[ -n "$extends" ]]; then
      ui::rule::list_extends_detailed "$extends" "$rules_dir"
      echo ""
    fi

  done <<< "$data"

  $found_any || {
    [[ -n "$filter_group" ]] && \
      log::wg_warning "No rules found in group: ${filter_group}" || \
      log::wg_warning "No rules found"
  }

  echo ""
}

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

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

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)     name="$2";          shift 2 ;;
      --no-peers) show_peers=false;   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")"

  # DNS display
  local dns_redirect resolved_dns dns_display
  dns_redirect=$(rule::get_own "$name" "dns_redirect")
  dns_redirect="${dns_redirect:-false}"
  resolved_dns=$(rule::get "$name" "dns_redirect")
  resolved_dns="${resolved_dns:-false}"

  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"

  local desc group
  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"

  printf "\n"
  if ui::rule::tree "$name"; then
    :
  else
    ui::rule::flat "$name"
    printf "\n"
  fi

  # Resolved view
  if $show_resolved; then
    ui::rule::section_header "Resolved (applied to peers)"
    printf "\n"
    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" ]] && net::print_entry "+" "$e"; done \
      <<< "$res_allow_ports"$'\n'"$res_allow_ips"
    while IFS= read -r e; do [[ -n "$e" ]] && net::print_entry "-" "$e"; done \
      <<< "$res_block_ips"$'\n'"$res_block_ports"
    printf "\n"
  fi

  # Peers
  $show_peers || return 0

  local peer_list=()
  mapfile -t peer_list < <(peers::with_rule "$name") || true
  local peer_count=${#peer_list[@]}
  ui::empty "$peer_count" && return 0

  local peer_word="peers"
  [[ "$peer_count" -eq 1 ]] && peer_word="peer"
  printf "\n  \033[0;37m── Peers (%s %s) \033[0m%s\n\n" \
    "$peer_count" "$peer_word" "$(printf '\033[0;37m─%.0s' {1..30})"

  for peer_name in "${peer_list[@]}"; do
    local ip
    ip=$(peers::get_ip "$peer_name")
    printf "    %-28s %s\n" "$peer_name" "$ip"
  done
  printf "\n"
  return 0
}

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

function cmd::rule::add() {
  local name="" desc="" group=""
  local extends=()
  local allow_ips=() block_ips=() block_ports=() allow_ports=()
  local block_services=() allow_services=()
  local dns_redirect=false 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 ;;
      --block-service) block_services+=("$2");  shift 2 ;;
      --allow-service) allow_services+=("$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

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

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

  for svc in "${block_services[@]}"; do
    while IFS= read -r resolved; do
      if [[ "$resolved" == *:*:* ]]; then
        block_ports+=("${resolved}")
      else
        block_ips+=("${resolved}/32")
      fi
    done < <(net::resolve "$svc")
  done

  for svc in "${allow_services[@]}"; do
    while IFS= read -r resolved; do
      if [[ "$resolved" == *:*:* ]]; then
        allow_ports+=("${resolved}")
      else
        allow_ips+=("${resolved}/32")
      fi
    done < <(net::resolve "$svc")
  done

  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")"

  [[ -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"

  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

  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

  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
# ============================================

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") || true
  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}"
}

# ============================================
# Assign / Unassign
# ============================================

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

  # Identity rule check
  local peer_identity
  
  peer_identity=$(peers::get_meta "$peer" "identity")
  if [[ -n "$peer_identity" ]]; then
    local identity_rules
    identity_rules=$(identity::rules "$peer_identity" 2>/dev/null)
    if echo "$identity_rules" | grep -qx "$name"; then
      log::error "Rule '${name}' is already applied to '${peer}' via identity '${peer_identity}' — cannot assign directly"
      return 1
    fi
  fi

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

  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}"
}

# ============================================
# Migrate
# ============================================

function cmd::rule::migrate() {
  log::section "Migrating peers to default rules"
  local count=0

  while IFS= read -r peer_name; do
    local existing
    existing=$(peers::get_meta "$peer_name" "rule")
    [[ -n "$existing" ]] && continue

    # Try to get default rule from subnet policy
    local peer_type subnet_name default_rule
    peer_type=$(peers::get_meta "$peer_name" "type")
    subnet_name=$(peers::get_meta "$peer_name" "subnet")
    default_rule=$(subnet::default_rule "$subnet_name" "$peer_type")
    [[ -z "$default_rule" ]] && continue

    rule::exists "$default_rule" || continue

    local ip
    ip=$(peers::get_ip "$peer_name")
    rule::apply "$default_rule" "$ip" "$peer_name"
    (( count++ )) || true
  done < <(peers::all)

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

# ============================================
# Reapply
# ============================================

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

  if $all; then
    log::section "Reapplying all rules"
    local count=0
    while IFS= read -r rule_file; do
      local rname
      rname=$(basename "$rule_file" .rule)
      local peer_list=()
      mapfile -t peer_list < <(peers::with_rule "$rname") || true
      [[ ${#peer_list[@]} -eq 0 ]] && continue
      rule::reapply_all "$rname"
      (( count++ )) || true
    done < <(find "$(ctx::rules)" -maxdepth 1 -name "*.rule")
    log::wg_success "Reapplied ${count} assignable rules"
    return 0
  fi

  [[ -z "$name" ]] && log::error "Missing --name or --all" && return 1
  rule::require_exists "$name" || return 1
  rule::reapply_all "$name"
  log::wg_success "Rule '${name}' reapplied"
}