# frozen_string_literal: true

module Gitlab
  module BackgroundMigration
    # This class populates the `finding_uuid` attribute for
    # the existing `vulnerability_feedback` records.
    class PopulateFindingUuidForVulnerabilityFeedback
      REPORT_TYPES = {
        sast: 0,
        dependency_scanning: 1,
        container_scanning: 2,
        dast: 3,
        secret_detection: 4,
        coverage_fuzzing: 5,
        api_fuzzing: 6
      }.freeze

      class VulnerabilityFeedback < ActiveRecord::Base # rubocop:disable Style/Documentation
        include EachBatch

        self.table_name = 'vulnerability_feedback'

        enum category: REPORT_TYPES

        scope :in_range, -> (start, stop) { where(id: start..stop) }
        scope :without_uuid, -> { where(finding_uuid: nil) }

        def self.load_vulnerability_findings
          all.to_a.tap { |collection| collection.each(&:vulnerability_finding) }
        end

        def set_finding_uuid
          return unless vulnerability_finding.present? && vulnerability_finding.primary_identifier.present?

          update_column(:finding_uuid, calculated_uuid)
        rescue StandardError => error
          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error)
        end

        def vulnerability_finding
          BatchLoader.for(finding_key).batch(replace_methods: false) do |finding_keys, loader|
            project_ids = finding_keys.map { |key| key[:project_id] }
            categories = finding_keys.map { |key| key[:category] }
            fingerprints = finding_keys.map { |key| key[:project_fingerprint] }

            findings = Finding.with_primary_identifier.where(
              project_id: project_ids.uniq,
              report_type: categories.uniq,
              project_fingerprint: fingerprints.uniq
            ).to_a

            finding_keys.each do |finding_key|
              loader.call(
                finding_key,
                findings.find { |f| finding_key == f.finding_key }
              )
            end
          end
        end

        private

        def calculated_uuid
          ::Security::VulnerabilityUUID.generate(
            report_type: category,
            primary_identifier_fingerprint: vulnerability_finding.primary_identifier.fingerprint,
            location_fingerprint: vulnerability_finding.location_fingerprint,
            project_id: project_id
          )
        end

        def finding_key
          {
            project_id: project_id,
            category: category,
            project_fingerprint: project_fingerprint
          }
        end
      end

      class Finding < ActiveRecord::Base # rubocop:disable Style/Documentation
        include ShaAttribute

        self.table_name = 'vulnerability_occurrences'

        sha_attribute :project_fingerprint
        sha_attribute :location_fingerprint

        belongs_to :primary_identifier, class_name: 'Gitlab::BackgroundMigration::PopulateFindingUuidForVulnerabilityFeedback::Identifier'

        enum report_type: REPORT_TYPES

        scope :with_primary_identifier, -> { includes(:primary_identifier) }

        def finding_key
          {
            project_id: project_id,
            category: report_type,
            project_fingerprint: project_fingerprint
          }
        end
      end

      class Identifier < ActiveRecord::Base # rubocop:disable Style/Documentation
        self.table_name = 'vulnerability_identifiers'
      end

      def perform(*range)
        feedback = VulnerabilityFeedback.without_uuid.in_range(*range).load_vulnerability_findings
        feedback.each(&:set_finding_uuid)

        log_info(feedback.count)
      end

      def log_info(feedback_count)
        ::Gitlab::BackgroundMigration::Logger.info(
          migrator: self.class.name,
          message: '`finding_uuid` attributes has been set',
          count: feedback_count
        )
      end
    end
  end
end