# frozen_string_literal: true

module EE
  module Gitlab
    module BackgroundMigration
      # This class creates/updates those project historical vulnerability statistics
      # that haven't been created nor initialized.
      module PopulateVulnerabilityHistoricalStatistics
        extend ::Gitlab::Utils::Override

        override :perform
        def perform(project_ids, retention_period = 90)
          project_ids.each do |project_id|
            upsert_vulnerability_historical_statistics(project_id, retention_period)
          rescue => e
            error_message("Error updating statistics for project #{project_id}: #{e.message}")
          end
        end

        private

        MAX_DAYS_IN_SINGLE_QUERY = 10

        EMPTY_STATISTIC = {
          total: 0,
          critical: 0,
          high: 0,
          medium: 0,
          low: 0,
          unknown: 0,
          info: 0
        }.freeze

        class Vulnerability < ActiveRecord::Base
          self.table_name = 'vulnerabilities'

          enum severity: { info: 1, unknown: 2, low: 4, medium: 5, high: 6, critical: 7 }
        end

        class VulnerabilityHistoricalStatistic < ActiveRecord::Base
          self.table_name = 'vulnerability_historical_statistics'

          enum letter_grade: { a: 0, b: 1, c: 2, d: 3, f: 4 }
        end

        def upsert_vulnerability_historical_statistics(project_id, retention_period)
          end_date = Date.today
          start_date = end_date - retention_period.days
          time_now = Time.current

          counts_by_day_and_severity_in_batches(project_id, start_date, end_date, of: MAX_DAYS_IN_SINGLE_QUERY)
            .map { |day, statistics| prepare_historical_statistic_attributes(project_id, day, statistics, time_now) }
            .then { |statistics| deduplicate_statistics(statistics) }
            .then { |statistics| VulnerabilityHistoricalStatistic.insert_all(statistics) }
        end

        def counts_by_day_and_severity_in_batches(project_id, start_date, end_date, of:)
          (start_date...end_date)
            .each_slice(of)
            .flat_map { |date_range| counts_by_day_and_severity(project_id, date_range.first, date_range.last) }
            .group_by(&:day)
        end

        def counts_by_day_and_severity(project_id, start_date, end_date)
          quoted_start_date = ActiveRecord::Base.connection.quote(start_date)
          quoted_end_date = ActiveRecord::Base.connection.quote(end_date)

          Vulnerability
            .where(project_id: project_id)
            .select('DATE(calendar.entry) AS day, severity, COUNT(*)')
            .from("generate_series(DATE #{quoted_start_date}, DATE #{quoted_end_date}, INTERVAL '1 day') as calendar(entry)")
            .joins('INNER JOIN vulnerabilities ON vulnerabilities.created_at <= calendar.entry')
            .where('(vulnerabilities.dismissed_at IS NULL OR vulnerabilities.dismissed_at > calendar.entry) AND (vulnerabilities.resolved_at IS NULL OR vulnerabilities.resolved_at > calendar.entry)')
            .group(:day, :severity)
        end

        def prepare_historical_statistic_attributes(project_id, day, statistics, time_now)
          severity_counts = statistics.map { |statistic| { statistic.severity.to_sym => statistic.count } }.inject(:merge)

          EMPTY_STATISTIC.merge(
            date: day,
            total: statistics.sum(&:count),
            letter_grade: letter_grade_for(severity_counts),
            created_at: time_now,
            updated_at: time_now,
            project_id: project_id,
            **severity_counts
          )
        end

        def deduplicate_statistics(statistics)
          statistics
            .sort_by { |statistic| statistic[:date] }
            .slice_when { |statistic_before, statistic| statistic_before.except(:date) != statistic.except(:date) }
            .map(&:first)
        end

        def letter_grade_for(statistic)
          if statistic[:critical].to_i > 0
            VulnerabilityHistoricalStatistic.letter_grades[:f]
          elsif statistic[:high].to_i > 0 || statistic[:unknown].to_i > 0
            VulnerabilityHistoricalStatistic.letter_grades[:d]
          elsif statistic[:medium].to_i > 0
            VulnerabilityHistoricalStatistic.letter_grades[:c]
          elsif statistic[:low].to_i > 0
            VulnerabilityHistoricalStatistic.letter_grades[:b]
          else
            VulnerabilityHistoricalStatistic.letter_grades[:a]
          end
        end

        def logger
          @logger ||= ::Gitlab::BackgroundMigration::Logger.build
        end

        def error_message(message)
          logger.error(message: "Vulnerability Historical Statistics Migration: #{message}")
        end
      end
    end
  end
end
