#!/usr/bin/env bash

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

function cmd::rule::on_load() {
  flag::register --name
  flag::register --desc
  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
}

# ============================================
# 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 and assigned peers
  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

Options for add/update:
  --name <name>               Rule name (e.g. guest, user, dev-01)
  --desc <description>        Human readable description
  --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 (update only)
  --remove-allow-port <entry> Remove allow port entry (update only)
  --remove-block-ip <ip>      Remove block IP entry (update only)
  --remove-block-port <entry> Remove block port entry (update only)

Options for assign/unassign:
  --name <rule>               Rule name
  --peer <peer>               Peer name (e.g. phone-nuno)
  --type <type>               Peer device type (optional)

Examples:
  wgctl rule list
  wgctl rule show --name guest
  wgctl rule add --name dev-01 --desc "Dev VM only" --allow-ip 10.0.0.50 --block-ip 10.0.0.0/24
  wgctl rule update --name user --block-port 10.0.0.100:8006:tcp
  wgctl rule update --name user --remove-block-ip 10.0.0.99
  wgctl rule assign --name dev-01 --peer laptop-nuno
  wgctl rule unassign --peer laptop-nuno --type laptop
  wgctl rule migrate
EOF
}

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

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

  case "$subcmd" in
    list|ls)              cmd::rule::list "$@"     ;;
    show)                 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 "$@"  ;;
    inspect)              cmd::rule::inspect "$@"  ;;
    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::list() {
  local rules_dir
  rules_dir="$(ctx::rules)"

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

  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"

  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 rules section header
    if [[ "$is_base" == "True" && "$printing_base" == "false" ]]; then
      local bdashes
      bdashes=$(printf '─%.0s' {1..74})
      printf "\n  \033[0;37m── Base Rules \033[0m%s\n" "$bdashes"
      printing_base=true
      current_group=""  # reset group tracking for base section
    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
        # Switching from grouped back to ungrouped
        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

    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 IMMEDIATELY after the rule row
    if [[ -n "$extends" ]]; then
      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
      printf "\n"  # blank line after rule with extends
    fi

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

  printf "\n"
}

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

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

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)   name="$2";        shift 2 ;;
      --peers)  show_peers=true;  shift   ;;
      --color)  color=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="$(ctx::rule::path "${name}.rule")"

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

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

  local desc dns_redirect
  desc=$(json::get "$rule_file" "desc")
  dns_redirect=$(json::get "$rule_file" "dns_redirect")

  printf "\n"
  ui::row "Description" "${desc:-—}"
  ui::row "DNS Redirect" "${dns_redirect:-false}"

  # Load all entries
  local allow_ports allow_ips block_ips block_ports
  allow_ports=$(json::get "$rule_file" "allow_ports")
  allow_ips=$(json::get "$rule_file" "allow_ips")
  block_ips=$(json::get "$rule_file" "block_ips")
  block_ports=$(json::get "$rule_file" "block_ports")

  # Allow section
  if [[ -n "$allow_ports" || -n "$allow_ips" ]]; then
    cmd::rule::_show_section "Allow" "green" "$color"
    cmd::rule::_show_entries "Ports" "+" "$allow_ports" "$color" "green"
    cmd::rule::_show_entries "IPs"   "+" "$allow_ips"   "$color" "green"
  fi

  # Block section
  if [[ -n "$block_ips" || -n "$block_ports" ]]; then
    cmd::rule::_show_section "Block" "red" "$color"
    cmd::rule::_show_entries "IPs"   "-" "$block_ips"   "$color" "red"
    cmd::rule::_show_entries "Ports" "-" "$block_ports" "$color" "red"
  fi

  if [[ -z "$allow_ports" && -z "$allow_ips" && -z "$block_ips" && -z "$block_ports" ]]; then
    printf "\n"
    ui::row "Access" "full (no restrictions)"
  fi

  # Peers section
  cmd::rule::_show_section "Peers" "white" false
  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"
}

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

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

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

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

#   local rule_file
#   rule_file="$(ctx::rule::path "${name}.rule")"

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

#   local desc dns_redirect
#   desc=$(json::get "$rule_file" "desc")
#   dns_redirect=$(json::get "$rule_file" "dns_redirect")

#   printf "\n"
#   ui::row "Description" "${desc:-—}"
#   ui::row "DNS Redirect" "${dns_redirect:-false}"

#   # Allow ports
#   local allow_ports
#   allow_ports=$(json::get "$rule_file" "allow_ports")
#   if [[ -n "$allow_ports" ]]; then
#     printf "\n  Allow Ports:\n"
#     ui::print_list "+" "$allow_ports"
#   fi

#   # Allow IPs
#   local allow_ips
#   allow_ips=$(json::get "$rule_file" "allow_ips")
#   if [[ -n "$allow_ips" ]]; then
#     printf "\n  Allow IPs:\n"
#     while IFS= read -r ip; do
#       printf "    + %s\n" "$ip"
#     done <<< "$allow_ips"
#   fi

#   # Block IPs
#   local block_ips
#   block_ips=$(json::get "$rule_file" "block_ips")
#   if [[ -n "$block_ips" ]]; then
#     printf "\n  Block IPs:\n"
#     while IFS= read -r ip; do
#       printf "    - %s\n" "$ip"
#     done <<< "$block_ips"
#   fi

#   # Block ports
#   local block_ports
#   block_ports=$(json::get "$rule_file" "block_ports")
#   if [[ -n "$block_ports" ]]; then
#     printf "\n  Block Ports:\n"
#     while IFS= read -r entry; do
#       printf "    - %s\n" "$entry"
#     done <<< "$block_ports"
#   fi

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

#   # Peer count — always shown
#   printf "\n  %-20s %s\n" "Assigned Peers:" "$peer_count"
#   printf "  %s\n" "$(printf '─%.0s' {1..40})"

#   # Peer details — only with --peers flag
#   if $show_peers && [[ $peer_count -gt 0 ]]; then
#     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=""

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) util::require_flag "--name" "${2:-}" || return 1; name="$2"; shift 2 ;;
      --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=""
  while IFS="|" read -r section key value; do
    [[ -z "$section" ]] && continue

    # Print section header when section changes
    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 to peers)" ;;
        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")

  printf "\n"
}

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

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

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

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

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

  local rule_file
  rule_file="$(ctx::rule::path "${name}.rule")"

  local allow_str block_str port_str

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

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

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

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

function cmd::rule::update() {
  local name="" desc=""
  local allow_ips=() block_ips=() block_ports=()
  local 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 ;;
      --allow-ip)          allow_ips+=("$2");        shift 2 ;;
      --block-ip)          block_ips+=("$2");        shift 2 ;;
      --block-port)        block_ports+=("$2");      shift 2 ;;
      --allow-port)        allow_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

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

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

  local rule_file
  rule_file="$(ctx::rule::path "${name}.rule")"

  # Update desc and dns_redirect
  [[ -n "$desc" ]] && json::set "$rule_file" "desc" "\"$desc\""
  [[ -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

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

  # Re-apply to all assigned peers
  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

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

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

  # Check for assigned peers

  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

    # Force: unassign from all peers
    for peer in "${peer_list[@]}"; do
      local ip
      ip=$(peers::get_ip "$peer")
      rule::unapply "$name" "$ip"
    done
  fi

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

# ============================================
# Assign
# ============================================

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

  if [[ -z "$name" || -z "$peer" ]]; then
    log::error "Missing required flags: --name and --peer"
    return 1
  fi

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

  # Support --type for peer resolution
  peer=$(peers::resolve_and_require "$peer" "$type") || return 1

  # Unapply existing rule first if any
  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}"
}

# ============================================
# Unassign
# ============================================

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

  if [[ -z "$peer" ]]; then
    log::error "Missing required flag: --peer"
    return 1
  fi

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

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

  # Write migration plan to temp file to avoid fd conflicts
  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"

  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
    
  cat "$tmp"

  rm -f "$tmp"
  log::wg_success "Migrated ${count} peers"
}

# ============================================
# Reapply Rule
# ============================================

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

function cmd::rule::_show_entries() {
  local label="${1:-}" prefix="${2:-}" entries="${3:-}"
  local use_color="${4:-false}" color="${5:-white}"
  [[ -z "$entries" ]] && return 0

  local color_code="" reset=""
  if $use_color; then
    case "$color" in
      green) color_code="\033[0;32m" ;;
      red)   color_code="\033[0;31m" ;;
    esac
    reset="\033[0m"
  fi

  printf "    %-8s" "${label}:"
  local first=true
  while IFS= read -r entry; do
    [[ -z "$entry" ]] && continue
    if $first; then
      printf "%b%s %s%b\n" "$color_code" "$prefix" "$entry" "$reset"
      first=false
    else
      printf "            %b%s %s%b\n" "$color_code" "$prefix" "$entry" "$reset"
    fi
  done <<< "$entries"
}