# frozen_string_literal: true

module Gitlab
  module Graphql
    module Aggregations
      module VulnerabilityStatistics
        class LazyAggregate
          attr_reader :vulnerable, :lazy_state

          def initialize(query_ctx, vulnerable)
            @vulnerable = vulnerable.respond_to?(:sync) ? vulnerable.sync : vulnerable

            # Initialize the loading state for this query,
            # or get the previously-initiated state
            @lazy_state = query_ctx[:lazy_vulnerability_statistics_aggregate] ||= {
              pending_vulnerables: Set.new,
              loaded_objects: {}
            }
            # Register this ID to be loaded later:
            @lazy_state[:pending_vulnerables] << vulnerable
          end

          # Return the loaded record, hitting the database if needed
          def execute
            # Check if the record was already loaded
            if @lazy_state[:pending_vulnerables].present?
              load_records_into_loaded_objects
            end

            @lazy_state[:loaded_objects][@vulnerable]
          end

          private

          def load_records_into_loaded_objects
            # The record hasn't been loaded yet, so
            # hit the database with all pending IDs to prevent N+1
            pending_vulnerables = @lazy_state[:pending_vulnerables].to_a
            grades = ::Vulnerabilities::ProjectsGrade.grades_for(pending_vulnerables)

            pending_vulnerables.each do |vulnerable|
              @lazy_state[:loaded_objects][vulnerable] = grades[vulnerable]
            end

            @lazy_state[:pending_vulnerables].clear
          end
        end
      end
    end
  end
end
