# Reboot OpenWiFi (openwifi_wss) APs in staggered batches, with live feedback.
#
# Uses AccessPoint#reboot -- the same path as the UI's Reboot button (reboot
# sent over the AP's WSS session, DeviceRebootEvent recorded).
#
# Batches are interleaved: AP i goes in batch (i mod BATCHES), with APs sorted
# by name, so each building loses only ~1/BATCHES of its APs at a time instead
# of going dark together. Staggering also spreads the rxg_sdan module download
# every AP does after boot, which a simultaneous restart turned into an outage.
#
# USAGE -- on the rXg, as root (csh-safe; settings are trailing arguments):
#
#   curl -o /tmp/ap-staggered-reboot.rb https://pre-alpha.henryhaller.com/ap-staggered-reboot.rb
#   cd /space/rxg/console
#   bin/rails runner /tmp/ap-staggered-reboot.rb                    # dry run: shows the plan
#   bin/rails runner /tmp/ap-staggered-reboot.rb APPLY=1            # do it
#   bin/rails runner /tmp/ap-staggered-reboot.rb APPLY=1 MODEL=186h # only matching models
#
#   APPLY=1        actually reboot (otherwise only print the plan)
#   BATCHES=10     number of batches (default 10)
#   INTERVAL=60    seconds between batches (default 60)
#   MODEL=text     only APs whose model contains text (case-insensitive)
#   ZONE=id        only one access point zone
#   LIMIT=n        cap the number of APs (try LIMIT=3 first)
#   WAIT=900       after the last batch, keep reporting until every rebooted AP
#                  has resubscribed or this many seconds pass (default 900)
#
# Ctrl-C stops dispatching; APs already told to reboot will still reboot.
ARGV.each do |arg|
  key, value = arg.split('=', 2)
  next if value.nil? || !key.match?(/\A[A-Z][A-Z0-9_]*\z/)
  ENV[key] ||= value
end

APPLY    = ENV['APPLY'] == '1'
BATCHES  = [(ENV['BATCHES'] || 10).to_i, 1].max
INTERVAL = (ENV['INTERVAL'] || 60).to_i
WAIT     = (ENV['WAIT'] || 900).to_i
$stdout.sync = true

def stamp
  Time.current.strftime('%H:%M:%S')
end

scope = AccessPoint.joins(:infrastructure_device)
                   .where(infrastructure_devices: { device: 'openwifi_wss' })
                   .where(online: true)
scope = scope.where(access_point_zone_id: ENV['ZONE']) if ENV['ZONE'].present?
scope = scope.where('access_points.model ILIKE ?', "%#{ENV['MODEL']}%") if ENV['MODEL'].present?
aps = scope.order('access_points.name').to_a
aps = aps.first(ENV['LIMIT'].to_i) if ENV['LIMIT'].present?
abort 'No online openwifi_wss APs matched.' if aps.empty?

batches = Array.new(BATCHES) { [] }
aps.each_with_index { |ap, i| batches[i % BATCHES] << ap }
batches.reject!(&:empty?)

puts "#{stamp} #{aps.size} online AP(s) in #{batches.size} batch(es) of ~#{batches.first.size}, " \
     "#{INTERVAL}s apart#{APPLY ? '' : '  [DRY RUN -- add APPLY=1 to reboot]'}"
aps.group_by { |a| a.model.presence || '(unknown)' }.each { |m, g| puts "    #{m}: #{g.size}" }

unless APPLY
  batches.each_with_index do |b, i|
    puts "  batch #{i + 1}: #{b.first(5).map(&:name).join(', ')}#{b.size > 5 ? ", ... (+#{b.size - 5})" : ''}"
  end
  exit 0
end

ids = aps.map(&:id)
rebooted_at = {}
failed = {}

# Progress comes from the WSS subscription timestamps: down = unsubscribed
# since its reboot was sent, back = subscribed again since then.
status = lambda do |label|
  rows = AccessPoint.where(id: ids).pluck(:id, :last_subscribed_at, :last_unsubscribed_at)
  connected = rows.count { |_, sub, unsub| sub && (unsub.nil? || sub > unsub) }
  back = Set.new
  went_down = 0
  rows.each do |id, sub, unsub|
    sent = rebooted_at[id]
    next unless sent
    back << id if sub && sub > sent
    went_down += 1 if (unsub && unsub > sent) || (sub && sub > sent)
  end
  puts format('%s %-14s connected %3d/%-3d | rebooted %3d: %3d went down, %3d back up | failed %d',
              stamp, label, connected, ids.size, rebooted_at.size, went_down, back.size, failed.size)
  back
end

batches.each_with_index do |batch, i|
  ok = 0
  batch.each do |ap|
    res = begin
      ap.reload.reboot
    rescue StandardError => e
      "#{e.class}: #{e.message}"
    end
    if res == true
      rebooted_at[ap.id] = Time.current
      ok += 1
    else
      failed[ap.id] = res.nil? ? 'reboot not supported / no client' : res.to_s
      puts "#{stamp}   ! #{ap.name}: #{failed[ap.id]}"
    end
  end
  puts "#{stamp} batch #{i + 1}/#{batches.size}: reboot sent to #{ok}/#{batch.size}"
  break if i == batches.size - 1

  # Report every 15 s while waiting for the next batch.
  deadline = Time.current + INTERVAL
  while Time.current < deadline
    sleep [15, (deadline - Time.current).ceil].min
    status.call("  waiting")
  end
end

puts "#{stamp} all batches sent; waiting up to #{WAIT}s for APs to come back"
deadline = Time.current + WAIT
back = Set.new
loop do
  back = status.call('  recovering')
  break if rebooted_at.keys.all? { |id| back.include?(id) } || Time.current >= deadline
  sleep 30
end

stragglers = rebooted_at.keys.reject { |id| back.include?(id) }
puts
puts "#{stamp} done: #{rebooted_at.size - stragglers.size}/#{rebooted_at.size} rebooted APs back (resubscribed), " \
     "#{stragglers.size} not back, #{failed.size} failed to dispatch"
AccessPoint.where(id: stragglers).order(:name).each { |ap| puts "    not back: #{ap.name} (#{ap.model}, id #{ap.id})" }
failed.each { |id, why| puts "    failed:   #{AccessPoint.find_by(id: id)&.name || id}: #{why}" }
puts 'Next: run sdan-fleet-check.rb (LIMIT=500 TIMEOUT=300) to confirm rxg_sdan loaded everywhere.'