#!/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
}

# ============================================
# 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 "$@"  ;;
    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 rules=("${rules_dir}"/*.rule)
  if [[ ! -f "${rules[0]}" ]]; then
    log::wg "No rules configured"
    return 0
  fi

  log::section "Firewall Rules"
  printf "\n  %-20s %-45s %-8s %-8s %s\n" \
    "NAME" "DESCRIPTION" "ALLOWS" "BLOCKS" "PEERS"
  printf "  %s\n" "$(printf '─%.0s' {1..95})"

  while IFS="|" read -r name desc n_allows n_blocks peer_count; do
    [[ -z "$name" ]] && continue
    local short_desc="${desc:0:43}"
    [[ ${#desc} -gt 43 ]] && short_desc="${short_desc}..."
    printf "  %-20s %-45s %-8s %-8s %s\n" \
      "$name" "${short_desc:-—}" "$n_allows" "$n_blocks" "${peer_count} peers"
  done < <(json::rule_list_data "$rules_dir" "$(ctx::meta)")

  printf "\n"
}

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

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  %-20s %s\n" "Description:" "${desc:-—}"
  printf "  %-20s %s\n" "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"
    while IFS= read -r entry; do
      printf "    + %s\n" "$entry"
    done <<< "$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"
}

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

  # Build JSON using json_helper
  python3 -c "
import json
rule = {
    'name': '${name}',
    'desc': '${desc}',
    'dns_redirect': $(${dns_redirect} && echo 'true' || echo 'false'),
    'allow_ips': [$(printf '"%s",' "${allow_ips[@]}" | sed 's/,$//')]  ,
    'block_ips': [$(printf '"%s",' "${block_ips[@]}" | sed 's/,$//')],
    'block_ports': [$(printf '"%s",' "${block_ports[@]}" | sed 's/,$//')]
}
with open('${rule_file}', 'w') as f:
    json.dump(rule, f, indent=2)
"

  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_count
  peer_count=$(peers::with_rule "$name" | grep -c . || echo 0)
  if [[ "$peer_count" -gt 0 ]]; then
    log::error "Rule '${name}' is assigned to ${peer_count} peer(s) — unassign first or use --force"
    if ! $force; then
      return 1
    fi
    # Force: unassign from all peers
    while IFS= read -r peer; do
      local ip
      ip=$(peers::get_ip "$peer")
      rule::unapply "$name" "$ip"
    done < <(peers::with_rule "$name")
  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

  # 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")
  if [[ -n "$existing_rule" ]]; then
    local ip
    ip=$(peers::get_ip "$peer")
    rule::unapply "$existing_rule" "$ip"
    log::wg "Removed existing rule '${existing_rule}' from: ${peer}"
  fi

  local ip
  ip=$(peers::get_ip "$peer")
  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"
  echo "DEBUG: lines count=${#lines[@]}"
for line in "${lines[@]}"; do
    echo "DEBUG: processing line=$line"
    IFS=" " read -r peer_name default_rule ip <<< "$line"
    rule::apply "$default_rule" "$ip" "$peer_name" </dev/null
    echo "DEBUG: after apply, count=$count"
    (( count++ )) || true
    echo "DEBUG: incremented count=$count"
  done
  echo "DEBUG: loop done, count=$count"
    
    echo "DEBUG: final count=$count"
    echo "DEBUG: tmp contents:"
    cat "$tmp"

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