# frozen_string_literal: true

require 'runners/crud/requirement'
require 'runners/unified/ambiguous_operations'
require 'runners/unified/client_side_encryption_operations'
require 'runners/unified/crud_operations'
require 'runners/unified/grid_fs_operations'
require 'runners/unified/ddl_operations'
require 'runners/unified/change_stream_operations'
require 'runners/unified/support_operations'
require 'runners/unified/thread_operations'
require 'runners/unified/search_index_operations'
require 'runners/unified/assertions'
require 'support/utils'
require 'support/crypt'

module Unified
  class Test
    include AmbiguousOperations
    include ClientSideEncryptionOperations
    include CrudOperations
    include GridFsOperations
    include DdlOperations
    include ChangeStreamOperations
    include SupportOperations
    include ThreadOperations
    include SearchIndexOperations
    include Assertions
    include RSpec::Core::Pending

    def initialize(spec, **opts)
      @spec = spec
      @entities = EntityMap.new
      @test_spec = UsingHash[@spec.fetch('test')]
      @description = @test_spec.use('description')
      @outcome = @test_spec.use('outcome')
      @expected_events = @test_spec.use('expectEvents')
      @expected_tracing_messages = @test_spec.use('expectTracingMessages')
      @expected_spans = @test_spec.use('expectSpans')
      @skip_reason = @test_spec.use('skipReason')
      if (req = @test_spec.use('runOnRequirements'))
        @reqs = req.map { |r| Mongo::CRUD::Requirement.new(r) }
      end
      if (req = @spec['group_runOnRequirements'])
        @group_reqs = req.map { |r| Mongo::CRUD::Requirement.new(r) }
      end
      if @spec['createEntities']
        mongoses = @spec['createEntities'].select do |spec|
          spec['client']
        end.map do |spec|
          spec['client']['useMultipleMongoses']
        end.compact.uniq
        @multiple_mongoses = mongoses.any? { |v| v }
      end
      @test_spec.freeze
      @subscribers = {}
      @tracers = {}
      @observe_sensitive = {}
      @options = opts
    end

    attr_reader :test_spec, :description, :outcome, :skip_reason, :reqs, :group_reqs, :options, :entities

    # Sentinel value used in unified spec KMS provider options to indicate that
    # the actual credential should be substituted from SpecConfig.
    KMS_PLACEHOLDER = { '$$placeholder' => 1 }.freeze

    # Descriptions of unified spec tests that are known to flake under load on
    # CI and benefit from being retried. See the corresponding JIRA tickets for
    # the underlying investigations.
    RETRY_PATTERNS = Regexp.union(
      /KMS/i,
      # RUBY-3809: CSOT deprecated-options "maxTimeMS is ignored if timeoutMS
      # is set" operations occasionally observe a drifted maxTimeMS on the
      # wire when the CI host is loaded.
      /maxTimeMS is ignored if timeoutMS is set/i,
      # RUBY-3810: sharded mongos-pin-auto "remain pinned after non-transient"
      # tests flake when the server injects its own transient transaction
      # condition around the failCommand-injected error.
      /remain pinned after non-transient/i
    ).freeze

    def retry?
      @description =~ RETRY_PATTERNS
    end

    def skip?
      !!@skip_reason
    end

    def require_multiple_mongoses?
      @multiple_mongoses == true
    end

    def require_single_mongos?
      @multiple_mongoses == false
    end

    def create_spec_entities
      return if @entities_created

      generate_entities(@spec['createEntities'])
    end

    def generate_entities(es)
      return if es.nil?

      es.each do |entity_spec|
        raise NotImplementedError, 'Entity must have exactly one key' unless entity_spec.keys.length == 1

        type, spec = entity_spec.first
        spec = UsingHash[spec]
        id = spec.use!('id')

        entity = case type
                 when 'client'
                   opts = if (smc_opts = spec.use('uriOptions'))
                            Mongo::URI::OptionsMapper.new.smc_to_ruby(smc_opts)
                          else
                            {}
                          end

                   # max_pool_size gets automatically set to 3 if not explicitly set by
                   # the test, therefore, if min_pool_size is set, make sure to set the
                   # max_pool_size as well to something greater.
                   if !opts.key?('max_pool_size') && (min_pool_size = opts[:min_pool_size])
                     opts[:max_pool_size] = min_pool_size + 3
                   end

                   if spec.use('useMultipleMongoses')
                     if (ClusterConfig.instance.topology == :sharded) && !(SpecConfig.instance.addresses.length > 1)
                       raise 'useMultipleMongoses requires more than one address in MONGODB_URI'
                     end
                   else
                     # If useMultipleMongoses isn't true, truncate the address
                     # list to the first address.
                     # This works OK in replica sets because the driver will discover
                     # the other set members, in standalone deployments because
                     # there is only one server, but changes behavior in
                     # sharded clusters compared to how the test suite is configured.
                     options[:single_address] = true
                   end

                   if (store_events = spec.use('storeEventsAsEntities'))
                     store_event_names = {}
                     store_events.each do |spec|
                       entity_name = spec['id']
                       event_names = spec['events']
                       event_names.each do |event_name|
                         store_event_names[event_name] = entity_name
                       end
                     end
                     store_event_names.values.uniq.each do |entity_name|
                       entities.set(:event_list, entity_name, [])
                     end
                     subscriber = StoringEventSubscriber.new do |payload|
                       if (entity_name = store_event_names[payload['name']])
                         entities.get(:event_list, entity_name) << payload
                       end
                     end
                     opts[:sdam_proc] = lambda do |client|
                       client.subscribe(Mongo::Monitoring::COMMAND, subscriber)
                       client.subscribe(Mongo::Monitoring::CONNECTION_POOL, subscriber)
                     end
                   end

                   if (server_api = spec.use('serverApi'))
                     server_api = ::Utils.underscore_hash(server_api)
                     opts[:server_api] = server_api
                   end

                   observe_events = spec.use('observeEvents')
                   subscriber = EventSubscriber.new
                   current_proc = opts[:sdam_proc]
                   opts[:sdam_proc] = lambda do |client|
                     current_proc.call(client) if current_proc
                     if (oe = observe_events)
                       oe.each do |event|
                         case event
                         when 'commandStartedEvent', 'commandSucceededEvent', 'commandFailedEvent'
                           unless client.send(:monitoring).subscribers[Mongo::Monitoring::COMMAND].include?(subscriber)
                             client.subscribe(Mongo::Monitoring::COMMAND, subscriber)
                           end
                           kind = event.sub('command', '').sub('Event', '').downcase.to_sym
                           subscriber.add_wanted_events(kind)
                           if (ignore_events = spec.use('ignoreCommandMonitoringEvents'))
                             subscriber.ignore_commands(ignore_events)
                           end
                         when /\A(?:pool|connection)/
                           unless client.send(:monitoring).subscribers[Mongo::Monitoring::CONNECTION_POOL]&.include?(subscriber)
                             client.subscribe(Mongo::Monitoring::CONNECTION_POOL, subscriber)
                           end
                           kind = event.sub('Event', '').gsub(/([A-Z])/) do
                             "_#{::Regexp.last_match(1)}"
                           end.sub('pool', 'Pool').downcase.to_sym
                           subscriber.add_wanted_events(kind)
                         when 'serverDescriptionChangedEvent'
                           unless client.send(:monitoring).subscribers[Mongo::Monitoring::SERVER_DESCRIPTION_CHANGED]&.include?(subscriber)
                             client.subscribe(Mongo::Monitoring::SERVER_DESCRIPTION_CHANGED, subscriber)
                           end
                           kind = event.sub('Event', '').gsub(/([A-Z])/) do
                             "_#{::Regexp.last_match(1)}"
                           end.downcase.to_sym
                           subscriber.add_wanted_events(kind)
                         when 'serverHeartbeatStartedEvent', 'serverHeartbeatSucceededEvent', 'serverHeartbeatFailedEvent'
                           unless client.send(:monitoring).subscribers[Mongo::Monitoring::SERVER_HEARTBEAT]&.include?(subscriber)
                             client.subscribe(Mongo::Monitoring::SERVER_HEARTBEAT, subscriber)
                           end
                           kind = event.sub('Event', '').gsub(/([A-Z])/) do
                             "_#{::Regexp.last_match(1)}"
                           end.downcase.to_sym
                           subscriber.add_wanted_events(kind)
                         else
                           raise NotImplementedError, "Unknown event #{event}"
                         end
                       end
                     end
                   end

                   observe_tracing_messages = spec.use('observeTracingMessages')
                   tracer = ::Tracing::Tracer.new
                   if observe_tracing_messages
                     opts[:tracing] = {
                       enabled: true,
                       tracer: tracer,
                     }
                     if observe_tracing_messages['enableCommandPayload']
                       # Set the maximum length of the query text to reasonably high
                       # value so that we can capture the full query text
                       opts[:tracing][:query_text_max_length] = 4096
                     end
                   end

                   await_min_pool_size_ms = spec.use('awaitMinPoolSizeMS')

                   if (auto_encrypt_opts = spec.use('autoEncryptOpts'))
                     opts.merge!(Utils.convert_client_options('autoEncryptOpts' => auto_encrypt_opts))
                   end

                   create_client(**opts).tap do |client|
                     @observe_sensitive[id] = spec.use('observeSensitiveCommands')
                     @subscribers[client] ||= subscriber
                     @tracers[client] ||= tracer

                     if await_min_pool_size_ms
                       min_pool_size = opts[:min_pool_size] || 0
                       if min_pool_size > 0
                         deadline = Mongo::Utils.monotonic_time + (await_min_pool_size_ms / 1000.0)
                         loop do
                           pools_ready = client.cluster.servers.all? do |server|
                             next true unless server.description.data_bearing?

                             begin
                               server.pool.size >= min_pool_size
                             rescue Mongo::Error::ServerNotUsable
                               false
                             end
                           end
                           break if pools_ready
                           if Mongo::Utils.monotonic_time > deadline
                             raise "Connection pool did not reach minPoolSize=#{min_pool_size} within #{await_min_pool_size_ms}ms"
                           end

                           sleep 0.01
                         end
                       end
                     end
                   end
                 when 'database'
                   client = entities.get(:client, spec.use!('client'))
                   opts = Utils.snakeize_hash(spec.use('databaseOptions') || {})
                               .merge(database: spec.use!('databaseName'))
                   if opts.key?(:read_preference)
                     opts[:read] = opts.delete(:read_preference)
                     if opts[:read].key?(:max_staleness_seconds)
                       opts[:read][:max_staleness] = opts[:read].delete(:max_staleness_seconds)
                     end
                   end
                   client.with(opts).database
                 when 'collection'
                   database = entities.get(:database, spec.use!('database'))
                   # TODO: verify
                   opts = Utils.snakeize_hash(spec.use('collectionOptions') || {})
                   if opts.key?(:read_preference)
                     opts[:read] = opts.delete(:read_preference)
                     if opts[:read].key?(:max_staleness_seconds)
                       opts[:read][:max_staleness] = opts[:read].delete(:max_staleness_seconds)
                     end
                   end
                   database[spec.use!('collectionName'), opts]
                 when 'bucket'
                   database = entities.get(:database, spec.use!('database'))
                   database.fs
                 when 'session'
                   client = entities.get(:client, spec.use!('client'))
                   opts = build_session_options(spec)

                   client.start_session(**opts).tap do |session|
                     session.advance_cluster_time(@cluster_time)
                   end
                 when 'clientEncryption'
                   client_encryption_opts = spec.use!('clientEncryptionOpts')
                   key_vault_client = entities.get(:client, client_encryption_opts['keyVaultClient'])
                   opts = {
                     key_vault_namespace: client_encryption_opts['keyVaultNamespace'],
                     kms_providers: client_encryption_opts['kmsProviders']
                                    .transform_keys { |k| Utils.underscore(k.to_s) }
                                    .transform_values do |provider_opts|
                                      provider_opts.transform_keys { |k| Utils.underscore(k.to_s) }
                                    end,
                     kms_tls_options: {
                       kmip: {
                         ssl_cert: SpecConfig.instance.fle_kmip_tls_certificate_key_file,
                         ssl_key: SpecConfig.instance.fle_kmip_tls_certificate_key_file,
                         ssl_ca_cert: SpecConfig.instance.fle_kmip_tls_ca_file
                       }
                     }
                   }
                   opts[:kms_providers] = opts[:kms_providers].map do |provider, options|
                     base_type = Mongo::Crypt::KMS.provider_base_type(provider).to_sym
                     converted_options = options.map do |key, value|
                       [ key, resolve_kms_provider_option(base_type, provider, key, value) ]
                     end.to_h
                     [ provider, converted_options ]
                   end.to_h

                   Mongo::ClientEncryption.new(
                     key_vault_client,
                     opts
                   )
                 when 'thread'
                   thread_context = ThreadContext.new
                   thread = Thread.new do
                     loop do
                       begin
                         op_spec = thread_context.operations.pop(true)
                         execute_operation(op_spec)
                       rescue ThreadError
                         # Queue is empty
                       end
                       break if thread_context.stop?

                       sleep 0.1
                     end
                   end
                   class << thread
                     attr_accessor :context
                   end
                   thread.context = thread_context
                   thread
                 else
                   raise NotImplementedError, "Unknown type #{type}"
                 end
        raise NotImplementedError, "Unhandled spec keys: #{spec}" unless spec.empty?

        entities.set(type.to_sym, id, entity)
      end
      @entities_created = true
    end

    # Builds the keyword options for Client#start_session from a session
    # entity spec. When sessionOptions contains snapshotTime, the value is the
    # name of a previously saved entity holding the actual BSON::Timestamp.
    def build_session_options(spec)
      smc_opts = spec.use('sessionOptions')
      return {} unless smc_opts

      opts = ::Utils.underscore_hash(smc_opts)
      if opts[:snapshot_time].is_a?(String)
        opts[:snapshot_time] = entities.get(:result, opts[:snapshot_time])
      end
      opts
    end

    def set_initial_data
      @spec['initialData']&.each do |entity_spec|
        spec = UsingHash[entity_spec]
        collection = root_authorized_client.with(write_concern: { w: :majority })
                                           .use(spec.use!('databaseName'))[spec.use!('collectionName')]
        collection.drop
        create_options = spec.use('createOptions') || {}
        docs = spec.use!('documents')
        begin
          collection.create(create_options)
        rescue Mongo::Error => e
          if e.is_a?(Mongo::Error::OperationFailure::Family) && (
              e.code == 48 || e.message =~ /collection already exists/
            )
            # Already exists
          else
            raise
          end
        end
        collection.insert_many(docs) if docs.any?
        raise NotImplementedError, "Unhandled spec keys: #{spec}" unless spec.empty?
      end

      # the cluster time is used to advance the cluster time of any
      # sessions created during this test.
      # -> see DRIVERS-2816
      @cluster_time = root_authorized_client.command(ping: 1).cluster_time
    end

    def run
      kill_sessions

      test_spec = UsingHash[self.test_spec]
      ops = test_spec.use!('operations')
      execute_operations(ops)
      raise NotImplementedError, "Unhandled spec keys: #{test_spec}" unless test_spec.empty?
    ensure
      disable_fail_points
    end

    def stop!
      @stop = true
    end

    def stop?
      !!@stop
    end

    def cleanup
      if $kill_transactions || true
        kill_sessions
        $kill_transactions = nil
      end

      disable_fail_points

      entities[:client]&.each do |_id, client|
        client.close
      end

      # Reset entity state so that retry_test re-runs (RUBY-3809, RUBY-3810)
      # build fresh clients/sessions instead of reusing the closed ones
      # produced by the previous attempt. Without this, retried transactional
      # tests hit TRANSACTION_ALREADY_IN_PROGRESS on the stale session.
      @entities = EntityMap.new
      @entities_created = false
    end

    private

    # Resolves a single KMS provider option value. If the value is the
    # placeholder sentinel, substitutes the appropriate credential from
    # SpecConfig. If the value is a base64-encoded local key string, decodes
    # it. Otherwise returns the value unchanged.
    def resolve_kms_provider_option(base_type, provider, key, value)
      if value == KMS_PLACEHOLDER
        resolved = case base_type
                   when :aws
                     case key
                     when :access_key_id
                       provider.to_s.end_with?(':name2') ? SpecConfig.instance.fle_aws_key2 : SpecConfig.instance.fle_aws_key
                     when :secret_access_key
                       provider.to_s.end_with?(':name2') ? SpecConfig.instance.fle_aws_secret2 : SpecConfig.instance.fle_aws_secret
                     end
                   when :azure
                     case key
                     when :tenant_id then SpecConfig.instance.fle_azure_tenant_id
                     when :client_id then SpecConfig.instance.fle_azure_client_id
                     when :client_secret then SpecConfig.instance.fle_azure_client_secret
                     end
                   when :gcp
                     case key
                     when :email then SpecConfig.instance.fle_gcp_email
                     when :private_key then SpecConfig.instance.fle_gcp_private_key
                     end
                   when :kmip
                     case key
                     when :endpoint then SpecConfig.instance.fle_kmip_endpoint
                     end
                   when :local
                     case key
                     when :key then Crypt::LOCAL_MASTER_KEY
                     end
                   end
        resolved || raise("No placeholder value configured for KMS provider #{provider}, key #{key}")
      elsif base_type == :local && key == :key && value.is_a?(String)
        Base64.strict_decode64(value)
      else
        value
      end
    end

    def execute_operations(ops, propagate_errors: false)
      ops.each do |op|
        execute_operation(op, propagate_errors: propagate_errors)
      end
    end

    def execute_operation(op, propagate_errors: false)
      use_all(op, 'operation', op) do |op|
        name = Utils.underscore(op.use!('name'))
        method_name = name
        method_name = "_#{name}" if name.to_s == 'loop'

        skip "Mongo Ruby Driver does not support #{name}" if %w[modify_collection list_index_names client_bulk_write].include?(name.to_s)

        if (expected_error = op.use('expectError'))
          begin
            unless respond_to?(method_name)
              raise Error::UnsupportedOperation, "Mongo Ruby Driver does not support #{name}"
            end

            public_send(method_name, op)
          rescue Mongo::Error, bson_error, Mongo::Auth::Unauthorized, ArgumentError => e
            if expected_error.use('isTimeoutError') && !e.is_a?(Mongo::Error::TimeoutError)
              raise e
              raise Error::ErrorMismatch, %,Expected TimeoutError ("isTimeoutError") but got #{e},
            end
            # isClientError doesn't actually mean a client error.
            # It means anything other than OperationFailure. DRIVERS-1799
            if expected_error.use('isClientError') && e.is_a?(Mongo::Error::OperationFailure::Family)
              raise Error::ErrorMismatch, %,Expected not OperationFailure ("isClientError") but got #{e},
            end
            if (code = expected_error.use('errorCode')) && !(e.code == code)
              raise Error::ErrorMismatch, "Expected #{code} code but had #{e.code}"
            end
            if (code_name = expected_error.use('errorCodeName')) && !(e.code_name == code_name)
              raise Error::ErrorMismatch, "Expected #{code_name} code name but had #{e.code_name}"
            end
            if (text = expected_error.use('errorContains')) && !e.to_s.include?(text)
              raise Error::ErrorMismatch, "Expected #{text} in the message but had #{e}"
            end

            if (labels = expected_error.use('errorLabelsContain'))
              labels.each do |label|
                unless e.label?(label)
                  raise Error::ErrorMismatch, "Expected error to contain label #{label} but it did not"
                end
              end
            end
            if (omit_labels = expected_error.use('errorLabelsOmit'))
              omit_labels.each do |label|
                raise Error::ErrorMismatch, "Expected error to not contain label #{label} but it did" if e.label?(label)
              end
            end
            if (error_response = expected_error.use('errorResponse'))
              assert_result_matches(e.document, error_response)
            end
            if (expected_result = expected_error.use('expectResult'))
              assert_result_matches(e.result, expected_result)
            # Important: this must be the last branch.
            elsif expected_error.use('isError')
              # Nothing but we consume the key.
            end
            raise NotImplementedError, "Unhandled keys: #{expected_error}" unless expected_error.empty?
            raise e if propagate_errors
          else
            raise Error::ErrorMismatch, 'Expected exception but none was raised'
          end
        elsif op.use('ignoreResultAndError')
          unless respond_to?(method_name)
            raise Error::UnsupportedOperation, "Mongo Ruby Driver does not support #{name}"
          end

          begin
            send(method_name, op)
          # We can possibly rescue more errors here, add as needed.
          rescue Mongo::Error
          end
        else
          unless respond_to?(method_name, true)
            raise Error::UnsupportedOperation, "Mongo Ruby Driver does not support #{name}"
          end

          result = send(method_name, op)
          if (expected_result = op.use('expectResult'))
            if result.nil? && expected_result.keys == [ '$$unsetOrMatches' ]
              return
            elsif result.nil? && !expected_result.empty?
              raise Error::ResultMismatch, "expected #{expected_result} but got nil"
            elsif expected_result.is_a?(Array)
              assert_documents_match(result, expected_result)
            else
              assert_result_matches(result, expected_result)
            end
            # expected_result.clear
          end
          if (save_entity = op.use('saveResultAsEntity'))
            entities.set(:result, save_entity, result)
          end
        end
      end
    end

    def use_sub(hash, key, &block)
      v = hash.use!(key)
      use_all(hash, key, v, &block)
    end

    def use_all(_hash, key, v)
      orig_v = v.dup
      (yield v).tap do
        raise NotImplementedError, "Unconsumed items for #{key}: #{v}\nOriginal hash: #{orig_v}" unless v.empty?
      end
    end

    def use_arguments(op, &block)
      if op.key?('arguments')
        use_sub(op, 'arguments', &block)
      else
        yield UsingHash.new
      end
    end

    def disable_fail_points
      return unless $disable_fail_points

      $disable_fail_points.each do |(fail_point_command, address)|
        client = ClusterTools.instance.direct_client(address,
                                                     database: 'admin')
        client.command(configureFailPoint: fail_point_command['configureFailPoint'],
                       mode: 'off')
      end
      $disable_fail_points = nil
    end

    def kill_sessions
      root_authorized_client.command(
        killAllSessions: []
      )
    rescue Mongo::Error::OperationFailure::Family => e
      if e.code == 11_601
        # operation was interrupted, ignore. SERVER-38335
      elsif e.code == 13
        # Unauthorized - e.g. when running in Atlas as part of
        # drivers-atlas-testing, ignore. SERVER-54216
      elsif e.code == 59
        # no such command (old server), ignore
      else
        raise
      end
    end

    def root_authorized_client
      @root_authorized_client ||= ClientRegistry.instance.global_client('root_authorized')
    end

    def create_client(**opts)
      args = case v = options[:client_args]
             when Array
               raise NotImplementedError, 'Client args array must have two elements' unless v.length == 2

               [ v.first, v.last.dup ]
             when String
               [ v, {} ]
             else
               addresses = SpecConfig.instance.addresses
               addresses = [ addresses.first ] if options[:single_address]
               [
                 addresses,
                 SpecConfig.instance.all_test_options,
               ]
             end
      args.last.update(
        max_read_retries: 0,
        max_write_retries: 0
      ).update(opts)
      Mongo::Client.new(*args)
    end

    # The error to rescue BSON tests for. If we still define
    # BSON::String::IllegalKey then we should rescue that particular error,
    # otherwise, rescue an arbitrary BSON::Error
    def bson_error
      if BSON::String.const_defined?(:IllegalKey)
        BSON::String.const_get(:IllegalKey)
      else
        BSON::Error
      end
    end
  end
end
