# frozen_string_literal: true

module Gitlab
  class Experiment
    module TestBehaviors
      autoload :Trackable, 'gitlab/experiment/test_behaviors/trackable.rb'
    end

    module RSpecHelpers
      def stub_experiments(experiments, times = nil)
        experiments.each { |experiment| wrapped_experiment(experiment, times) }
      end

      def wrapped_experiment(experiment, times = nil, expected = false, &block)
        klass, experiment_name, variant_name = *experiment_details(experiment)
        base_klass = Configuration.base_class.constantize

        # Set expectations on experiment classes so we can and_wrap_original with more specific args
        experiment_klasses = base_klass.descendants.reject { |k| k == klass }
        experiment_klasses.push(base_klass).each do |k|
          allow(k).to receive(:new).and_call_original
        end

        receiver = receive(:new)

        # Be specific for BaseClass calls
        receiver = receiver.with(experiment_name, any_args) if experiment_name && klass == base_klass

        receiver.exactly(times).times if times

        # Set expectations on experiment class of interest
        allow_or_expect_klass = expected ? expect(klass) : allow(klass)
        allow_or_expect_klass.to receiver.and_wrap_original do |method, *original_args, &original_block|
          method.call(*original_args).tap do |e|
            # Stub internal methods before calling the original_block
            allow(e).to receive(:enabled?).and_return(true)

            if variant_name == true # passing true allows the rollout to do its job
              allow(e).to receive(:experiment_group?).and_return(true)
            else
              allow(e).to receive(:resolve_variant_name).and_return(variant_name.to_s)
            end

            # Stub/set expectations before calling the original_block
            yield e if block

            original_block.call(e) if original_block.present?
          end
        end
      end

      private

      def experiment_details(experiment)
        if experiment.is_a?(Symbol)
          experiment_name = experiment
          variant_name = nil
        end

        experiment_name, variant_name = *experiment if experiment.is_a?(Array)

        base_klass = Configuration.base_class.constantize
        variant_name = experiment.variant.name if experiment.is_a?(base_klass)

        if experiment.class.name.nil? # Anonymous class instance
          klass = experiment.class
        elsif experiment.instance_of?(Class) # Class level stubbing, eg. "MyExperiment"
          klass = experiment
        else
          experiment_name ||= experiment.instance_variable_get(:@name)
          klass = base_klass.constantize(experiment_name)
        end

        if experiment_name && klass == base_klass
          experiment_name = experiment_name.to_sym

          # For experiment names like: "group/experiment-name"
          experiment_name = experiment_name.to_s if experiment_name.inspect.include?('"')
        end

        [klass, experiment_name, variant_name]
      end
    end

    module RSpecMatchers
      extend RSpec::Matchers::DSL

      def require_experiment(experiment, matcher_name, classes: false)
        klass = experiment.instance_of?(Class) ? experiment : experiment.class
        unless klass <= Gitlab::Experiment
          raise(
            ArgumentError,
            "#{matcher_name} matcher is limited to experiment instances#{classes ? ' and classes' : ''}"
          )
        end

        if experiment == klass && !classes
          raise ArgumentError, "#{matcher_name} matcher requires an instance of an experiment"
        end

        experiment != klass
      end

      matcher :exclude do |context|
        ivar = :'@excluded'

        match do |experiment|
          require_experiment(experiment, 'exclude')
          experiment.context(context)

          experiment.instance_variable_set(ivar, nil)
          !experiment.run_callbacks(:exclusion_check) { :not_excluded }
        end

        failure_message do
          %(expected #{context} to be excluded)
        end

        failure_message_when_negated do
          %(expected #{context} not to be excluded)
        end
      end

      matcher :segment do |context|
        ivar = :'@variant_name'

        match do |experiment|
          require_experiment(experiment, 'segment')
          experiment.context(context)

          experiment.instance_variable_set(ivar, nil)
          experiment.run_callbacks(:segmentation_check)

          @actual = experiment.instance_variable_get(ivar)
          @expected ? @actual.to_s == @expected.to_s : @actual.present?
        end

        chain :into do |expected|
          raise ArgumentError, 'variant name must be provided' if expected.blank?

          @expected = expected.to_s
        end

        failure_message do
          %(expected #{context} to be segmented#{message_details})
        end

        failure_message_when_negated do
          %(expected #{context} not to be segmented#{message_details})
        end

        def message_details
          message = ''
          message += %( into variant\n    expected variant: #{@expected}) if @expected
          message += %(\n      actual variant: #{@actual}) if @actual
          message
        end
      end

      matcher :track do |event, *event_args|
        match do |experiment|
          expect_tracking_on(experiment, false, event, *event_args)
        end

        match_when_negated do |experiment|
          expect_tracking_on(experiment, true, event, *event_args)
        end

        chain :for do |expected_variant|
          raise ArgumentError, 'variant name must be provided' if expected.blank?

          @expected_variant = expected_variant.to_s
        end

        chain(:with_context) { |expected_context| @expected_context = expected_context }

        chain(:on_next_instance) { @on_next_instance = true }

        def expect_tracking_on(experiment, negated, event, *event_args)
          klass = experiment.instance_of?(Class) ? experiment : experiment.class
          unless klass <= Gitlab::Experiment
            raise(
              ArgumentError,
              "track matcher is limited to experiment instances and classes"
            )
          end

          expectations = proc do |e|
            @experiment = e
            allow(e).to receive(:track).and_call_original

            if negated
              expect(e).not_to receive(:track).with(*[event, *event_args])
            else
              if @expected_variant
                expect(@experiment.variant.name).to eq(@expected_variant), failure_message(:variant, event)
              end

              if @expected_context
                expect(@experiment.context.value).to include(@expected_context), failure_message(:context, event)
              end

              expect(e).to receive(:track).with(*[event, *event_args]).and_call_original
            end
          end

          if experiment.instance_of?(Class) || @on_next_instance
            wrapped_experiment(experiment, nil, true) { |e| expectations.call(e) }
          else
            expectations.call(experiment)
          end
        end

        def failure_message(failure_type, event)
          case failure_type
          when :variant
            <<~MESSAGE.strip
              expected #{@experiment.inspect} to have tracked #{event.inspect} for variant
                  expected variant: #{@expected_variant}
                    actual variant: #{@experiment.variant.name}
            MESSAGE
          when :context
            <<~MESSAGE.strip
              expected #{@experiment.inspect} to have tracked #{event.inspect} with context
                  expected context: #{@expected_context}
                    actual context: #{@experiment.context.value}
            MESSAGE
          end
        end
      end
    end
  end
end

RSpec.configure do |config|
  config.include Gitlab::Experiment::RSpecHelpers
  config.include Gitlab::Experiment::Dsl

  config.before(:each, :experiment) do
    RequestStore.clear!

    if defined?(Gitlab::Experiment::TestBehaviors::TrackedStructure)
      Gitlab::Experiment::TestBehaviors::TrackedStructure.reset!
    end
  end

  config.include Gitlab::Experiment::RSpecMatchers, :experiment
  config.define_derived_metadata(file_path: Regexp.new('/spec/experiments/')) do |metadata|
    metadata[:type] = :experiment
  end
end
