# Copyright 2024 - 2025 Block, Inc.
#
# Use of this source code is governed by an MIT-style
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.
#
# frozen_string_literal: true

require "elastic_graph/admin/rake_tasks"
require "elastic_graph/local/docker_runner"
require "elastic_graph/schema_definition/rake_tasks"
require "net/http"
require "rake/tasklib"
require "shellwords"
require "uri"

module ElasticGraph
  # Provides support for developing and running ElasticGraph applications locally.
  module Local
    # Defines tasks for local development. These tasks include:
    #
    # - Running OpenSearch and/or Elasticsearch locally (`(elasticsearch|opensearch):[env]:(boot|daemon|halt)`)
    # - Managing schema artifacts (`schema_artifacts:(check|dump)`)
    # - Configuring OpenSearch/Elasticsearch locally (`clusters:configure:(dry_run|perform)`)
    # - Indexing fake data (`index_fake_data:[type]`)
    # - Booting an ElasticGraph application locally (`boot_locally`)
    #
    # @note All tasks (besides the `schema_artifacts` tasks) require `docker` and `docker-compose` to be available on your machine.
    class RakeTasks < ::Rake::TaskLib
      # When enabled, ElasticGraph will configure the index mappings so that the datastore indexes a `_size` field in each index document.
      # ElasticGraph itself does not do anything with this field, but it will be available for your use in any direct queries (e.g. via
      # Kibana).
      #
      # Defaults to `false` since it requires a plugin.
      #
      # @note Enabling this requires the [mapper-size plugin](https://www.elastic.co/guide/en/elasticsearch/plugins/8.15/mapper-size.html)
      #   to be installed on your datastore cluster. You are responsible for ensuring that is installed if you enable this feature. If you
      #   enable this and the plugin is not installed, you will get errors!
      #
      # @return [Boolean] whether or not the `_size` field should be indexed on each indexed type
      #
      # @example Enable indexing document sizes
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.index_document_sizes = true
      #   end
      #
      # @dynamic index_document_sizes, index_document_sizes=
      attr_accessor :index_document_sizes

      # The form of names for schema elements (fields, arguments, directives) generated by ElasticGraph, either `:snake_case` or
      # `:camelCase`. For example, if set to `:camelCase`, ElasticGraph will generate a `groupedBy` field, but if set to `:snake_case`,
      # ElasticGraph will generate a `grouped_by` field.
      #
      # Defaults to `:camelCase` since most GraphQL schemas use that casing.
      #
      # @return [:camelCase, :snake_case] which form to use
      #
      # @example Use `snake_case` names instead of `camelCase`
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.schema_element_name_form = :snake_case
      #   end
      #
      # @dynamic schema_element_name_form, schema_element_name_form=
      attr_accessor :schema_element_name_form

      # Overrides for specific names of schema elements (fields, arguments, directives) generated by ElasticGraph. For example, to rename
      # the `gt` filter field to `greaterThan`, set to `{gt: "greaterThan"}`.
      #
      # Defaults to an empty hash.
      #
      # @return [Hash<Symbol, String>] overrides for specific field, argument, or directive names
      #
      # @example Spell out comparison operators instead of using shortened forms
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.schema_element_name_overrides = {
      #       gt: "greaterThan",
      #       gte: "greaterThanOrEqualTo",
      #       lt: "lessThan",
      #       lte: "lessThanOrEqualTo"
      #     }
      #   end
      #
      # @dynamic schema_element_name_overrides, schema_element_name_overrides=
      attr_accessor :schema_element_name_overrides

      # Overrides for the naming formats used by ElasticGraph for derived GraphQL type names. For example, to use `Metrics` instead of
      # `AggregatedValues` as the suffix for the generated types supporting getting aggregated metrid values, set to
      # `{AggregatedValues: "%{base}Metrics"}`. See {SchemaDefinition::SchemaElements::TypeNamer::DEFAULT_FORMATS} for the available
      # formats.
      #
      # Defaults to an empty hash.
      #
      # @example Change the `AggregatedValues` type suffix to `Metrics`
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.derived_type_name_formats = {AggregatedValues: "%{base}Metrics"}
      #   end
      #
      # @dynamic derived_type_name_formats, derived_type_name_formats=
      attr_accessor :derived_type_name_formats

      # Overrides for the names of specific GraphQL types. For example, to rename the `JsonSafeLong` scalar to `BigInt`, set to
      # `{JsonSafeLong: "BigInt}`.
      #
      # Defaults to an empty hash.
      #
      # @return [Hash<Symbol, String>] overrides for specific type names
      #
      # @example Rename `JsonSafeLong` to `BigInt`
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.type_name_overrides = {JsonSafeLong: "BigInt"}
      #   end
      #
      # @dynamic type_name_overrides, type_name_overrides=
      attr_accessor :type_name_overrides

      # Overrides for the names of specific GraphQL enum values for specific enum types. For example, to rename the `DayOfWeek.MONDAY`
      # enum to `DayOfWeek.MON`, set to `{DayOfWeek: {MONDAY: "MON"}}`.
      #
      # Defaults to an empty hash.
      #
      # @return [Hash<Symbol, Hash<Symbol, String>>] overrides for the names of specific enum values for specific enum types
      #
      # @example Shorten the names of the `DayOfWeek` enum values
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.enum_value_overrides_by_type = {
      #       DayOfWeek: {
      #         MONDAY: "MON",
      #         TUESDAY: "TUE",
      #         WEDNESDAY: "WED",
      #         THURSDAY: "THU",
      #         FRIDAY: "FRI",
      #         SATURDAY: "SAT",
      #         SUNDAY: "SUN"
      #       }
      #     }
      #   end
      #
      # @dynamic enum_value_overrides_by_type, enum_value_overrides_by_type=
      attr_accessor :enum_value_overrides_by_type

      # List of Ruby modules to extend onto the {SchemaDefinition::API} instance. Designed to support ElasticGraph extensions (such as
      # {Apollo::SchemaDefinition::APIExtension}). Defaults to an empty list.
      #
      # @return [Array<Module>] list of extension modules
      #
      # @example Use `elasticgraph-apollo`
      #   require "elastic_graph/apollo/schema_definition/api_extension"
      #
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.schema_definition_extension_modules = [ElasticGraph::Apollo::SchemaDefinition::APIExtension]
      #   end
      #
      # @example Extension that defines a `@since` directive and offers a `since` API on fields
      #   module SinceExtension
      #     # `self.extended` is a standard Ruby hook that gets called when a module is extended onto an object.
      #     # The argument is the object the module was extended onto (a `SchemaDefinition::API` instance in this case).
      #     def self.extended(api)
      #       # Define our `@since` directive
      #       api.raw_sdl "directive @since(date: Date!) on FIELD_DEFINITION"
      #
      #       # In order to hook into fields, extend the `SchemaDefinition::Factory` with a module. The factory is used
      #       # for creation of all schema definition objects.
      #       api.factory.extend FactoryExtension
      #     end
      #
      #     module FactoryExtension
      #       # Hook into the creation of all `SchemaDefinition::Field` objects so that we can extend each field
      #       # instance with our `FieldExtension` module.
      #       def new_field(*args, **options)
      #         super(*args, **options) do |field|
      #           field.extend FieldExtension
      #           yield field if block_given?
      #         end
      #       end
      #     end
      #
      #     # Offer a `f.since date` API on fields.
      #     module FieldExtension
      #       def since(date)
      #         directive "since", date: date
      #       end
      #     end
      #   end
      #
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.schema_definition_extension_modules = [SinceExtension]
      #   end
      #
      # @dynamic schema_definition_extension_modules, schema_definition_extension_modules=
      attr_accessor :schema_definition_extension_modules

      # Whether or not to enforce the requirement that the JSON schema version is incremented every time
      # dumping the JSON schemas results in a changed artifact. Defaults to `true`.
      #
      # @note Generally speaking, you will want this to be `true` for any ElasticGraph application that is in
      #    production as the versioning of JSON schemas is what supports safe schema evolution as it allows
      #    ElasticGraph to identify which version of the JSON schema the publishing system was operating on
      #    when it published an event.
      #
      #    It can be useful to set it to `false` before your application is in production, as you do not want
      #    to be forced to bump the version after every single schema change while you are building an initial
      #    prototype.
      #
      # @return [Boolean] whether to require `json_schema_version` to be incremented on changes that impact `json_schemas.yaml`
      # @see SchemaDefinition::API#json_schema_version
      #
      # @example Disable enforcement during initial prototyping
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     # TODO: remove this once we're past the prototyping stage
      #     tasks.enforce_json_schema_version = false
      #   end
      #
      # @dynamic enforce_json_schema_version, enforce_json_schema_version=
      attr_accessor :enforce_json_schema_version

      # List of Elasticsearch versions you want to be able to boot. Rake tasks will be defined for each version to support booting and
      # halting Elasticsearch locally. If the configuration of `local_config_yaml` only configures `opensearch` as a cluster backend,
      # will default to an empty array. Otherwise, defaults to the versions of Elasticsearch that are exercised by the ElasticGraph test suite, as
      # defined by `lib/elastic_graph/local/tested_datastore_versions.yaml`:
      #
      # {include:file:elasticgraph-local/lib/elastic_graph/local/tested_datastore_versions.yaml}
      #
      # @return [Array<String>] list of Elasticsearch versions
      # @see #opensearch_versions
      #
      # @example Configure the Elasticsearch versions available to boot locally
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.elasticsearch_versions = ["8.16.1"]
      #   end
      #
      # @dynamic elasticsearch_versions, elasticsearch_versions=
      attr_accessor :elasticsearch_versions

      # List of OpenSearch versions you want to be able to boot. Rake tasks will be defined for each version to support booting and
      # halting OpenSearch locally. If the configuration of `local_config_yaml` only configures `elasticsearch` as a cluster backend,
      # will default to an empty array. Otherwise, defaults to the versions of OpenSearch that are exercised by the ElasticGraph test suite, as
      # defined by `lib/elastic_graph/local/tested_datastore_versions.yaml`:
      #
      # {include:file:elasticgraph-local/lib/elastic_graph/local/tested_datastore_versions.yaml}
      #
      # @return [Array<String>] list of OpenSearch versions
      # @see #elasticsearch_versions
      #
      # @example Configure the OpenSearch versions available to boot locally
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.opensearch_versions = ["2.18.0"]
      #   end
      #
      # @dynamic opensearch_versions, opensearch_versions=
      attr_accessor :opensearch_versions

      # Hash mapping environments (e.g. `:test`, `:dev`, etc) to port numbers for use when booting Elasticsearch or OpenSearch. The hash
      # automatically includes an entry for the `:local` environment, using a port number extracted from `local_config_yaml`.
      #
      # @note When booting Elasticsearch/OpenSearch, Kibana (or its OpenSearch equivalent, "OpenSearch Dashboards") will also get booted,
      #   selecting the port by adding `10000` to the configured port.
      #
      # @return [Hash<Symbol, Integer>] mapping from environment name to port number
      #
      # @example Define what port to use to boot the datastore for the `:test` environment
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.env_port_mapping = {test: 9999}
      #   end
      #
      # @dynamic env_port_mapping, env_port_mapping=
      attr_accessor :env_port_mapping

      # IO for printing output (defaults to stdout).
      #
      # @return [IO] IO object used for printing output.
      #
      # @dynamic output, output=
      attr_accessor :output

      # Maximum time (in seconds) to wait for the datastore to boot when booting it as a daemon. Defaults to 120.
      #
      # @return [Integer] maximum time in seconds to wait when booting Elasticsearch/OpenSearch as a daemon
      #
      # @dynamic daemon_timeout, daemon_timeout=
      attr_accessor :daemon_timeout

      # Offset we add to a port number for the UI (e.g. Kibana or OpenSearch Dashboards).
      #
      # Example: if Elasticsearch/OpenSearch is running on port 9876, the UI for it will run on port 19876.
      UI_PORT_OFFSET = 10_000

      # As per https://en.wikipedia.org/wiki/Registered_port, valid user port numbers are 1024 to 49151, but
      # with our UI offset we need to truncate the range further.
      #
      # @private
      VALID_PORT_RANGE = 1024..(49151 - UI_PORT_OFFSET)

      # Register a callback for use when indexing a batch fake data. An `index_fake_data:[type]` rake task will be generated for each
      # registered callback.
      #
      # @param type [Symbol] type of data batch. Can be the name of a GraphQL type or any other name you want to give a batch of fake data
      # @yieldreturn [Array<Hash<String, Object>>] generated data
      #
      # @example Register a callback to generate fake `campaigns` data
      #   ElasticGraph::Local::RakeTasks.new(
      #     local_config_yaml: "config/settings/local.yaml",
      #     path_to_schema: "config/schema.rb"
      #   ) do |tasks|
      #     tasks.define_fake_data_batch_for :campaigns do
      #       FactoryBot.build_list(:campaigns)
      #     end
      #   end
      def define_fake_data_batch_for(type, &block)
        @fake_data_batch_generator_by_type[type] = block
      end

      # @note This method uses keyword args for all required arguments. Optional task settings are instead specified using the block.
      # @param local_config_yaml [String, Pathname] path to the settings YAML file for the local/development environment
      # @param path_to_schema [String, Pathname] path to the Ruby schema definition file--either the only file that defines the schema
      #   (using `ElasticGraph.define_schema`) or the "main" schema definition file, which loads other files which further define parts of
      #   the schema.
      # @yield [RakeTasks] instance for further configuration
      # @yieldreturn [void]
      def initialize(local_config_yaml:, path_to_schema:)
        @local_config_yaml = local_config_yaml.to_s

        self.index_document_sizes = false
        self.schema_element_name_form = :camelCase
        self.schema_element_name_overrides = {}
        self.derived_type_name_formats = {}
        self.type_name_overrides = {}
        self.enum_value_overrides_by_type = {}
        self.schema_definition_extension_modules = []
        self.enforce_json_schema_version = true
        self.env_port_mapping = {}
        self.output = $stdout
        self.daemon_timeout = 300

        datastore_versions = ::YAML.load_file("#{__dir__}/tested_datastore_versions.yaml")
        self.elasticsearch_versions = local_cluster_backends.include?("elasticsearch") ? datastore_versions.fetch("elasticsearch") : []
        self.opensearch_versions = local_cluster_backends.include?("opensearch") ? datastore_versions.fetch("opensearch") : []

        @fake_data_batch_generator_by_type = {}

        yield self if block_given?

        # Default the local port from the local_config_yaml file.
        local_port = local_datastore_url[/.*:(\d+)$/, 1].then { |port_str| Integer(port_str) }
        self.env_port_mapping = {"local" => local_port}.merge(env_port_mapping || {})
        if (invalid_port_mapping = env_port_mapping.reject { |env, port| VALID_PORT_RANGE.cover?(port) }).any?
          raise "`env_port_mapping` has invalid ports: #{invalid_port_mapping.inspect}. Valid ports must be in the #{VALID_PORT_RANGE} range."
        end

        # Load admin and schema def rake tasks...
        Admin::RakeTasks.from_yaml_file(local_config_yaml, output: output)
        SchemaDefinition::RakeTasks.new(
          index_document_sizes: index_document_sizes,
          path_to_schema: path_to_schema,
          schema_artifacts_directory: local_config.fetch("schema_artifacts").fetch("directory"),
          schema_element_name_form: schema_element_name_form,
          schema_element_name_overrides: schema_element_name_overrides,
          derived_type_name_formats: derived_type_name_formats,
          type_name_overrides: type_name_overrides,
          enum_value_overrides_by_type: enum_value_overrides_by_type,
          extension_modules: schema_definition_extension_modules,
          enforce_json_schema_version: enforce_json_schema_version,
          output: output
        )

        # ...then define a bunch of our own.
        define_docker_tasks("Elasticsearch", "Kibana", elasticsearch_versions, /license \[[^\]]+\] mode \[[^\]]+\] - valid/)
        define_docker_tasks("OpenSearch", "OpenSearch Dashboards", opensearch_versions, /o\.o\.n\.Node.+started/)
        define_other_tasks
      end

      private

      def define_docker_tasks(description, ui_variant, versions, ready_log_line)
        variant = description.downcase.to_sym
        namespace variant do
          env_port_mapping.each do |env, port|
            namespace env do
              versions.each do |version|
                namespace version do
                  define_docker_tasks_for_version(description, variant, ui_variant, port: port, version: version, env: env, ready_log_line: ready_log_line)
                end
              end

              if (max_version = versions.max_by { |v| Gem::Version.create(v) })
                define_docker_tasks_for_version(description, variant, ui_variant, port: port, version: max_version, env: env, ready_log_line: ready_log_line)
              end
            end
          end
        end
      end

      def define_docker_tasks_for_version(description, variant, ui_variant, port:, version:, env:, ready_log_line:)
        ui_port = port + UI_PORT_OFFSET

        docker_runner = DockerRunner.new(
          variant,
          port: port,
          ui_port: port + UI_PORT_OFFSET,
          version: version,
          env: env,
          ready_log_line: ready_log_line,
          output: output,
          daemon_timeout: daemon_timeout
        )

        desc "Boots #{description} #{version} for the #{env} environment on port #{port} (and #{ui_variant} on port #{ui_port})"
        task(:boot) { docker_runner.boot }

        desc "Boots #{description} #{version} as a background daemon for the #{env} environment on port #{port} (and #{ui_variant} on port #{ui_port})"
        task(:daemon) do |t|
          docker_runner.boot_as_daemon(halt_command: "rake #{t.name.sub(/:\w+\z/, ":halt")}")
        end

        desc "Halts the #{description} #{version} daemon for the #{env} environment"
        task(:halt) { docker_runner.halt }
      end

      def define_other_tasks
        index_fake_data_tasks = @fake_data_batch_generator_by_type.keys.map do |type|
          "index_fake_data:#{type}"
        end

        datastore_to_boot =
          if elasticsearch_versions.empty? && opensearch_versions.empty?
            raise "Both `elasticsearch_versions` and `opensearch_versions` are empty, but we need at least one of them to have a version in order to provide the boot tasks."
          elsif elasticsearch_versions.empty?
            "OpenSearch"
          else
            "Elasticsearch"
          end

        desc "Boots ElasticGraph locally from scratch: boots #{datastore_to_boot}, configures it, indexes fake data, and boots GraphiQL"
        task :boot_locally, [:port, :rackup_args, :no_open] => ["#{datastore_to_boot.downcase}:local:daemon", *index_fake_data_tasks, "boot_graphiql"]

        desc "Boots ElasticGraph locally with the GraphiQL UI, and opens it in a browser."
        task :boot_graphiql, [:port, :rackup_args, :no_open] => :ensure_datastore_ready_for_indexing_and_querying do |task, args|
          args.with_defaults(port: 9000, rackup_args: "", no_open: false)
          port = args.fetch(:port)

          # :nocov: -- we can't test `open` behavior through a test
          unless args.fetch(:no_open)
            fork do
              sleep 3 # give the app a bit of time to boot before we try to open it.
              sh "open http://localhost:#{port}/"
            end
          end
          # :nocov:

          sh "ELASTICGRAPH_YAML_FILE=#{@local_config_yaml.shellescape} bundle exec rackup #{::File.join(__dir__.to_s, "config.ru").shellescape} --port #{port} #{args.fetch(:rackup_args)}"
        end

        namespace :index_fake_data do
          @fake_data_batch_generator_by_type.each do |type, generator|
            desc "Indexes num_batches of #{type} fake data into the local datastore"
            task type, [:num_batches] => :ensure_datastore_ready_for_indexing_and_querying do |task, args|
              require "elastic_graph/local/local_indexer"
              args.with_defaults(num_batches: 1)
              LocalIndexer.new(@local_config_yaml, generator, output: output).index_fake_data(Integer(args[:num_batches]))
            end
          end
        end

        task :ensure_local_datastore_running do
          datastore_status_code = begin
            ::Net::HTTP.get_response(URI(local_datastore_url)).code
          rescue ::SystemCallError, ::EOFError
            "500"
          end

          unless datastore_status_code == "200"
            ::Rake::Task["#{datastore_to_boot.downcase}:local:daemon"].invoke
          end
        end

        task ensure_datastore_ready_for_indexing_and_querying: [
          :ensure_local_datastore_running,
          "schema_artifacts:dump",
          "clusters:configure:perform"
        ]

        namespace "clusters:configure" do
          %i[dry_run perform].each do |subtask|
            desc "(after first dumping the schema artifacts)"
            task subtask => [:ensure_local_datastore_running, "schema_artifacts:dump"]
          end
        end
      end

      def local_datastore_url
        @local_datastore_url ||= local_config
          .fetch("datastore")
          .fetch("clusters")
          .fetch("main")
          .fetch("url")
      end

      def local_cluster_backends
        @local_cluster_backends ||= local_config
          .fetch("datastore")
          .fetch("clusters")
          .values
          .map { |v| v.fetch("backend") }
          .to_set
      end

      def local_config
        @local_config ||= ::YAML.safe_load_file(@local_config_yaml, aliases: true)
      end
    end
  end
end
