#!/usr/bin/env bash

# ============================================
# Rule File Parsing
# ============================================

function rule::exists() {
  local name="$1"
  [[ -f "$(ctx::rule::path "${name}.rule")" ]]
}

function rule::require_exists() {
  local name="$1"
  if ! rule::exists "$name"; then
    log::error "Rule not found: ${name}"
    return 1
  fi
}

function rule::get() {
  local name="$1" key="$2"
  json::get "$(ctx::rule::path "${name}.rule")" "$key"
}

function rule::get_all() {
  local name="$1"
  local rule_file
  rule_file="$(ctx::rule::path "${name}.rule")"
  cat "$rule_file"
}

function rule::is_applied() {
  local rule_name="$1"
  local client_ip="$2"

  # Try first block_ports entry
  local first_port
  first_port=$(rule::get "$rule_name" "block_ports" | head -1)
  if [[ -n "$first_port" ]]; then
    local target port proto
    IFS=":" read -r target port proto <<< "$first_port"
    proto="${proto:-tcp}"
    iptables -C FORWARD -s "$client_ip" -d "$target" -p "$proto" --dport "$port" -j DROP 2>/dev/null
    return $?
  fi

  # Fall back to first block_ips entry
  local first_ip
  first_ip=$(rule::get "$rule_name" "block_ips" | head -1)
  if [[ -n "$first_ip" ]]; then
    iptables -C FORWARD -s "$client_ip" -d "$first_ip" -j DROP 2>/dev/null
    return $?
  fi

  # No rules to check (admin rule) — check allow_ports
  local first_allow
  first_allow=$(rule::get "$rule_name" "allow_ports" | head -1)
  if [[ -n "$first_allow" ]]; then
    local target port proto
    IFS=":" read -r target port proto <<< "$first_allow"
    proto="${proto:-tcp}"
    iptables -C FORWARD -s "$client_ip" -d "$target" -p "$proto" --dport "$port" -j ACCEPT 2>/dev/null
    return $?
  fi

  # Empty rule (admin) — never "applied" in iptables sense
  return 1
}

# ============================================
# Rule Application
# ============================================

function rule::apply() {
  local rule_name="${1:?rule_name required}"
  local client_ip="${2:?client_ip required}"
  local peer_name="${3:-}"  # optional, avoids find_by_ip call

  rule::require_exists "$rule_name" || return 1

  # Use provided peer_name or look it up
  if [[ -z "$peer_name" ]]; then
    peer_name=$(peers::find_by_ip "$client_ip")
  fi

  # Check if already applied
  if rule::is_applied "$rule_name" "$client_ip"; then
    log::wg "Rule '${rule_name}' already applied to: ${client_ip}"
    
    # Still update meta even if rules exist
    if [[ -n "$peer_name" ]]; then
      peers::set_meta "$peer_name" "rule" "$rule_name"
    fi

    return 0
  fi

  # Check if already applied
  local peer_name
  peer_name=$(peers::find_by_ip "$client_ip")
  if [[ -n "$peer_name" ]]; then
    # Check if already applied via iptables
    if rule::is_applied "$rule_name" "$client_ip"; then
      log::wg "Rule '${rule_name}' already applied to: ${client_ip}"
      return 0
    fi
  fi

  # Process block_ips
  while IFS= read -r block_ip; do
    [[ -z "$block_ip" ]] && continue
    fw::block_ip "$client_ip" "$block_ip"
  done < <(rule::get "$rule_name" "block_ips")

  # Process block_ports
  while IFS= read -r entry; do
    [[ -z "$entry" ]] && continue
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    fw::block_port "$client_ip" "$target" "$port" "$proto"
  done < <(rule::get "$rule_name" "block_ports")

  # Process allow_ips (inserted before blocks)
  while IFS= read -r allow_ip; do
    [[ -z "$allow_ip" ]] && continue
    fw::allow_ip "$client_ip" "$allow_ip"
  done < <(rule::get "$rule_name" "allow_ips")

  # allow_ports (inserted last = highest priority)
  while IFS= read -r entry; do
    [[ -z "$entry" ]] && continue
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    fw::allow_port "$client_ip" "$target" "$port" "$proto"
  done < <(rule::get "$rule_name" "allow_ports")

  # Persist rule assignment in meta
  if [[ -n "$peer_name" ]]; then
    peers::set_meta "$peer_name" "rule" "$rule_name"
  fi

  local dns_redirect
  dns_redirect=$(rule::get "$rule_name" "dns_redirect")

  if [[ "$dns_redirect" == "true" ]]; then
    local peer_subnet
    peer_subnet=$(peers::get_ip "$peer_name" | cut -d'.' -f1-3)
    # Only apply if not already in PREROUTING
    if ! iptables -t nat -C PREROUTING -i wg0 -s "${peer_subnet}.0/24" -p udp --dport 53 \
      -j DNAT --to-destination "$(config::dns):53" 2>/dev/null; then
      rule::apply_dns_redirect "${peer_subnet}.0/24"
      log::debug "dns_redirect: applied for ${peer_subnet}.0/24"
    else
      log::debug "dns_redirect: already applied for ${peer_subnet}.0/24"
    fi
  fi

  log::debug "Applied rule '${rule_name}' to: ${client_ip}"
}

function rule::unapply() {
  local rule_name="$1"
  local client_ip="$2"

  rule::require_exists "$rule_name" || return 1

  local peer_name
  peer_name=$(peers::find_by_ip "$client_ip")

  local subtype
  subtype=$(peers::get_meta "$peer_name" "subtype")

  local subnet
  subnet=$(config::subnet_for "$subtype")

  # Remove allow_ports first (reverse order of apply)
  while IFS= read -r entry; do
    [[ -z "$entry" ]] && continue
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    fw::unallow_port "$client_ip" "$target" "$port" "$proto"
  done < <(rule::get "$rule_name" "allow_ports")

  # Remove allow_ips
  while IFS= read -r allow_ip; do
    [[ -z "$allow_ip" ]] && continue
    fw::unallow_ip "$client_ip" "$allow_ip"
  done < <(rule::get "$rule_name" "allow_ips")

  # Remove block_ports
  while IFS= read -r entry; do
    [[ -z "$entry" ]] && continue
    local target port proto
    IFS=":" read -r target port proto <<< "$entry"
    proto="${proto:-tcp}"
    fw::unblock_port "$client_ip" "$target" "$port" "$proto"
  done < <(rule::get "$rule_name" "block_ports")

  # Remove block_ips
  while IFS= read -r block_ip; do
    [[ -z "$block_ip" ]] && continue
    fw::unblock_ip "$client_ip" "$block_ip"
  done < <(rule::get "$rule_name" "block_ips")

  # Remove DNS redirect if applicable
  local dns_redirect
  dns_redirect=$(rule::get "$rule_name" "dns_redirect")
  if [[ "$dns_redirect" == "true" ]]; then
    if [[ -z "$subtype" ]]; then
      local peer_type
      peer_type=$(peers::get_type "$peer_name") || true
      [[ -z "$peer_type" ]] && peer_type="phone"  # fallback
      subnet=$(config::subnet_for "$peer_type")
    fi

    rule::remove_dns_redirect "${subnet}.0/24"
  fi

  # Clear rule from meta
  if [[ -n "$peer_name" ]]; then
    peers::set_meta "$peer_name" "rule" ""
  fi

  log::debug "Removed rule '${rule_name}' from: ${client_ip}"
}

function rule::reapply_all() {
  local rule_name="$1"
  rule::require_exists "$rule_name" || return 1

  local peers=()
  mapfile -t peers < <(peers::with_rule "$rule_name")
  
  [[ ${#peers[@]} -eq 0 ]] && return 0

  local count=0
  for peer_name in "${peers[@]}"; do
    local client_ip
    client_ip=$(peers::get_ip "$peer_name")
    [[ -z "$client_ip" ]] && continue
    rule::unapply "$rule_name" "$client_ip"
    rule::apply "$rule_name" "$client_ip" "$peer_name"
    (( count++ )) || true
  done

  log::wg_success "Rule '${rule_name}' re-applied to ${count} peers"
}

function rule::restore_all() {
  while IFS= read -r peer_name; do
    local rule_name
    rule_name=$(peers::get_meta "$peer_name" "rule")
    [[ -z "$rule_name" ]] && continue
    
    if ! rule::exists "$rule_name"; then
      log::wg_warning "Rule '${rule_name}' not found for peer '${peer_name}', skipping"
      continue
    fi
    
    local client_ip
    client_ip=$(peers::get_ip "$peer_name")
    [[ -z "$client_ip" ]] && continue
    
    rule::apply "$rule_name" "$client_ip"
  done < <(peers::all)
  log::wg "Rules restored for all peers"
}

# ============================================
# Guest DNS Redirect (rule-level feature)
# ============================================

function rule::apply_dns_redirect() {
  local client_subnet="$1"
  local dns
  dns="$(config::dns)"

  iptables -t nat -A PREROUTING -i wg0 -s "$client_subnet" -p udp --dport 53 \
    ! -d "$dns" -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4
  iptables -t nat -A PREROUTING -i wg0 -s "$client_subnet" -p udp --dport 53 \
    -j DNAT --to-destination "${dns}:53"
  iptables -t nat -A PREROUTING -i wg0 -s "$client_subnet" -p tcp --dport 53 \
    -j DNAT --to-destination "${dns}:53"
}

function rule::remove_dns_redirect() {
  local client_subnet="$1"
  local dns
  dns="$(config::dns)"

  iptables -t nat -D PREROUTING -i wg0 -s "$client_subnet" -p udp --dport 53 \
    ! -d "$dns" -j LOG --log-prefix "wgctl-dns-redirect: " --log-level 4 2>/dev/null || true
  iptables -t nat -D PREROUTING -i wg0 -s "$client_subnet" -p udp --dport 53 \
    -j DNAT --to-destination "${dns}:53" 2>/dev/null || true
  iptables -t nat -D PREROUTING -i wg0 -s "$client_subnet" -p tcp --dport 53 \
    -j DNAT --to-destination "${dns}:53" 2>/dev/null || true
}