# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

require 'contrast/components/logger'
require 'contrast/components/assess'

module Contrast
  module Utils
    module Assess
      # EventLimitUtils is used to check and validate the number of source, propagation, or trigger events collected
      # during the reporting time frame
      module EventLimitUtils
        include Contrast::Components::Logger::InstanceMethods
        # Checks to see if the event limit for the policy type has been met or exceeded
        # @param policy [Contrast::Agent::Patching::Policy::MethodPolicy,
        #   Contrast::Agent::Patching::Policy::TriggerNode] method to check for event limit
        def event_limit? policy
          return false unless (context = Contrast::Agent::REQUEST_TRACKER.current)

          if policy.source_node
            max =  ::Contrast::ASSESS.max_context_source_events
            return at_limit?(policy, context.source_event_count, max, context)
          end

          if policy.propagation_node
            max = ::Contrast::ASSESS.max_propagation_events
            return at_limit?(policy, context.propagation_event_count, max, context)
          end

          false # policy does not have limit
        end

        def event_limit_for_rule? rule_id # rubocop:disable Metrics/AbcSize
          return false unless (context = Contrast::Agent::REQUEST_TRACKER.current)

          saved_request_ids = rule_counts.keys.map { |k| k.to_s.split('_')[1].to_i }

          # if we passed the threshold and we actually have records for that request - wipe them
          if saved_request_ids.uniq.include?(context.request.__id__)
            restore_defaults
            threshold_time_limit
          end

          # if we have recorded rule counts, but none of them are for the current request_id
          # eventually we can try and play with the time_limit_threshold -> DEFAULT_MAX_RULE_TIME_THRESHOLD
          if !rule_counts.empty? && !saved_request_ids.include?(context.request.__id__)
            restore_defaults
            threshold_time_limit
          end

          rule_key = "#{ rule_id }_#{ context.request.__id__ }"
          rule_counts[rule_key] += 1
          rule_counts[rule_key] >= ::Contrast::ASSESS.max_rule_reported
        end

        # Increments the event count for the type of event that is being tracked
        #
        # @param node [Contrast::Agent::Assess::Policy::PolicyNode] policy to increment
        def increment_event_count node
          return unless (context = Contrast::Agent::REQUEST_TRACKER.current)

          context.source_event_count += 1 if node.cs__is_a?(Contrast::Agent::Assess::Policy::SourceNode)
          context.propagation_event_count += 1 if node.cs__is_a?(Contrast::Agent::Assess::Policy::PropagationNode)
        end

        private

        # helper method to check limit and log when necessary
        def at_limit? method_policy, current_count, event_max, context
          if current_count == event_max
            return if event_limit_counts.key?(get_event_limit_key(method_policy, context))

            logger.warn('Event Limit Reached:',
                        {
                            count: current_count,
                            max: event_max,
                            policy: method_policy.method_name,
                            node: method_policy
                        })
            # increment to be over count for logging purposes
            increment_event_count(method_policy)
            increment_event_limit_logs(method_policy, context)

            return true
          elsif current_count > event_max
            increment_event_count(method_policy)
            return if event_limit_counts.key?(get_event_limit_key(method_policy, context))

            # increment to be over count for logging purposes
            logger.warn('Event Limit Exceeded:',
                        {
                            count: current_count,
                            policy: method_policy.method_name,
                            node: method_policy
                        })
            increment_event_limit_logs(method_policy, context)
            return true
          end
          false
        end

        def rule_counts
          @_rule_counts ||= Hash.new { |h, k| h[k] = 0 }
        end

        def event_limit_counts
          @_event_limit_counts ||= Hash.new { |h, k| h[k] = 0 }
        end

        def get_event_limit_key method_policy, context
          "#{ method_policy.method_name }_#{ context.request.__id__ }"
        end

        def increment_event_limit_logs method_policy, context
          event_limit_counter = get_event_limit_key(method_policy, context)

          event_limit_counts[event_limit_counter] += 1
        end

        # the time threshold for which to track rule counts resets when now >= threshold_time_limit
        # @return [Integer]
        def threshold_time_limit
          @_threshold_time_limit ||= Contrast::Utils::Timer.now_ms + (::Contrast::ASSESS.time_limit_threshold || 0)
        end

        # @return nil
        def restore_defaults
          @_rule_counts = nil
          @_threshold_time_limit = nil
        end
      end
    end
  end
end