class AddEventsCountToNetworkRanges < ActiveRecord::Migration[8.1]
  def up
    # Add the column with default value
    add_column :network_ranges, :events_count, :integer, null: false, default: 0

    # Add index for faster queries
    add_index :network_ranges, :events_count

    # Create trigger function to update counter cache
    execute <<-SQL
      CREATE OR REPLACE FUNCTION update_network_range_events_count()
      RETURNS TRIGGER AS $$
      BEGIN
        -- Update all network ranges that contain the IP address
        UPDATE network_ranges
        SET events_count = events_count +
          CASE
            WHEN TG_OP = 'INSERT' THEN 1
            WHEN TG_OP = 'DELETE' THEN -1
            ELSE 0
          END
        WHERE network >>= NEW.ip_address::inet;

        RETURN COALESCE(NEW, OLD);
      END;
      $$ LANGUAGE plpgsql;
    SQL

    # Create triggers for events table
    execute <<-SQL
      CREATE TRIGGER update_network_ranges_events_count_after_insert
      AFTER INSERT ON events
      FOR EACH ROW
      EXECUTE FUNCTION update_network_range_events_count();
    SQL

    execute <<-SQL
      CREATE TRIGGER update_network_ranges_events_count_after_delete
      AFTER DELETE ON events
      FOR EACH ROW
      EXECUTE FUNCTION update_network_range_events_count();
    SQL

    # Backfill existing counts
    execute <<-SQL
      UPDATE network_ranges
      SET events_count = (
        SELECT COUNT(*)
        FROM events
        WHERE events.ip_address <<= network_ranges.network
      );
    SQL
  end

  def down
    # Drop triggers first
    execute <<-SQL
      DROP TRIGGER IF EXISTS update_network_ranges_events_count_after_insert ON events;
      DROP TRIGGER IF EXISTS update_network_ranges_events_count_after_delete ON events;
      DROP FUNCTION IF EXISTS update_network_range_events_count();
    SQL

    # Remove column and index
    remove_index :network_ranges, :events_count
    remove_column :network_ranges, :events_count
  end
end
