# frozen_string_literal: true

module Gitlab
  module Graphql
    module Aggregations
      module VulnerabilityStatistics
        class LazyAggregate
          include ::Gitlab::Graphql::Deferred

          attr_reader :vulnerable, :lazy_state

          def initialize(query_ctx, vulnerable, include_subgroups: false)
            @vulnerable = vulnerable.respond_to?(:sync) ? vulnerable.sync : vulnerable
            @include_subgroups = include_subgroups
            # Initialize the loading state for this query,
            # or get the previously-initiated state
            @lazy_state = query_ctx[:lazy_vulnerability_statistics_aggregate] ||= {
              pending_vulnerables: { true => Set.new, false => Set.new },
              loaded_objects: { true => {}, false => {} }
            }
            # Register this ID to be loaded later:
            @lazy_state[:pending_vulnerables][@include_subgroups] << 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][@include_subgroups].present?
              load_records_into_loaded_objects
            end

            @lazy_state[:loaded_objects][@include_subgroups][@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
            @lazy_state[:pending_vulnerables].each do |include_subgroups, pending_vulnerables|
              next if pending_vulnerables.blank?

              grades = ::Vulnerabilities::ProjectsGrade.grades_for(pending_vulnerables, include_subgroups: include_subgroups)

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

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