# frozen_string_literal: true

require "graphql/dataloader/null_dataloader"
require "graphql/dataloader/request"
require "graphql/dataloader/request_all"
require "graphql/dataloader/source"
require "graphql/dataloader/active_record_association_source"
require "graphql/dataloader/active_record_source"

module GraphQL
  # This plugin supports Fiber-based concurrency, along with {GraphQL::Dataloader::Source}.
  #
  # @example Installing Dataloader
  #
  #   class MySchema < GraphQL::Schema
  #     use GraphQL::Dataloader
  #   end
  #
  # @example Waiting for batch-loaded data in a GraphQL field
  #
  #   field :team, Types::Team, null: true
  #
  #   def team
  #     dataloader.with(Sources::Record, Team).load(object.team_id)
  #   end
  #
  class Dataloader
    class << self
      attr_accessor :default_nonblocking, :default_fiber_limit
    end

    def self.use(schema, nonblocking: nil, fiber_limit: nil)
      dataloader_class = if nonblocking
        warn("`nonblocking: true` is deprecated from `GraphQL::Dataloader`, please use `GraphQL::Dataloader::AsyncDataloader` instead. Docs: https://graphql-ruby.org/dataloader/async_dataloader.")
        Class.new(self) { self.default_nonblocking = true }
      else
        self
      end

      if fiber_limit
        dataloader_class = Class.new(dataloader_class)
        dataloader_class.default_fiber_limit = fiber_limit
      end

      schema.dataloader_class = dataloader_class
    end

    # Call the block with a Dataloader instance,
    # then run all enqueued jobs and return the result of the block.
    def self.with_dataloading(&block)
      dataloader = self.new
      result = nil
      dataloader.append_job {
        result = block.call(dataloader)
      }
      dataloader.run
      result
    end

    def initialize(nonblocking: self.class.default_nonblocking, fiber_limit: self.class.default_fiber_limit)
      @source_cache = Hash.new { |h, k| h[k] = {} }.compare_by_identity
      @pending_source_set = Set.new.compare_by_identity
      @pending_sources = []
      @pending_jobs = []
      if !nonblocking.nil?
        @nonblocking = nonblocking
      end
      @fiber_limit = fiber_limit
      @lazies_at_depth = Hash.new { |h, k| h[k] = [] }
    end

    # @return [Integer, nil]
    attr_reader :fiber_limit

    def nonblocking?
      @nonblocking
    end

    # This is called before the fiber is spawned, from the parent context (i.e. from
    # the thread or fiber that it is scheduled from).
    #
    # @return [Hash<Symbol, Object>] Current fiber-local variables
    def get_fiber_variables
      fiber_vars = {}
      Thread.current.keys.each do |fiber_var_key|
        fiber_vars[fiber_var_key] = Thread.current[fiber_var_key]
      end
      fiber_vars
    end

    # Set up the fiber variables in a new fiber.
    #
    # This is called within the fiber, right after it is spawned.
    #
    # @param vars [Hash<Symbol, Object>] Fiber-local variables from {get_fiber_variables}
    # @return [void]
    def set_fiber_variables(vars)
      vars.each { |k, v| Thread.current[k] = v }
      nil
    end

    # This method is called when Dataloader is finished using a fiber.
    # Use it to perform any cleanup, such as releasing database connections (if required manually)
    def cleanup_fiber
    end

    # Get a Source instance from this dataloader, for calling `.load(...)` or `.request(...)` on.
    #
    # @param source_class [Class<GraphQL::Dataloader::Source]
    # @param batch_parameters [Array<Object>]
    # @return [GraphQL::Dataloader::Source] An instance of {source_class}, initialized with `self, *batch_parameters`,
    #   and cached for the lifetime of this {Multiplex}.
    if (RUBY_ENGINE == "ruby" && RUBY_VERSION < "3") || RUBY_ENGINE == "truffleruby" # truffle-ruby wasn't doing well with the implementation below
      def with(source_class, *batch_args)
        batch_key = source_class.batch_key_for(*batch_args)
        @source_cache[source_class][batch_key] ||= begin
          source = source_class.new(*batch_args)
          source.setup(self)
          source
        end
      end
    else
      def with(source_class, *batch_args, **batch_kwargs)
        batch_key = source_class.batch_key_for(*batch_args, **batch_kwargs)
        @source_cache[source_class][batch_key] ||= begin
          source = source_class.new(*batch_args, **batch_kwargs)
          source.setup(self)
          source
        end
      end
    end
    # Tell the dataloader that this fiber is waiting for data.
    #
    # Dataloader will resume the fiber after the requested data has been loaded (by another Fiber).
    #
    # @return [void]
    def yield(source = Fiber[:__graphql_current_dataloader_source])
      trace = Fiber[:__graphql_current_multiplex]&.current_trace
      trace&.dataloader_fiber_yield(source)
      Fiber.yield
      trace&.dataloader_fiber_resume(source)
      nil
    end

    # @api private Nothing to see here
    def append_job(callable = nil, &job)
      # Given a block, queue it up to be worked through when `#run` is called.
      # (If the dataloader is already running, then a Fiber will pick this up later.)
      @pending_jobs.push(callable || job)
      nil
    end

    # @api private
    def queue_pending_source(source)
      if @pending_source_set.add?(source)
        @pending_sources << source
      end
      nil
    end

    # Clear any already-loaded objects from {Source} caches
    # @return [void]
    def clear_cache
      @source_cache.each do |_source_class, batched_sources|
        batched_sources.each_value(&:clear_cache)
      end
      nil
    end

    # Use a self-contained queue for the work in the block.
    def run_isolated
      prev_queue = @pending_jobs
      prev_pending_keys = {}
      prev_lazies_at_depth = @lazies_at_depth
      @lazies_at_depth = @lazies_at_depth.dup.clear
      # Clear pending loads but keep already-cached records
      # in case they are useful to the given block.
      @source_cache.each do |source_class, batched_sources|
        batched_sources.each do |batch_args, batched_source_instance|
          if batched_source_instance.pending?
            prev_pending_keys[batched_source_instance] = batched_source_instance.pending.dup
            batched_source_instance.pending.clear
          end
        end
      end

      @pending_jobs = []
      res = nil
      # Make sure the block is inside a Fiber, so it can `Fiber.yield`
      append_job {
        res = yield
      }
      run
      res
    ensure
      @pending_jobs = prev_queue
      @lazies_at_depth = prev_lazies_at_depth
      prev_pending_keys.each do |source_instance, pending|
        pending.each do |key, value|
          next if source_instance.results.key?(key)

          queue_pending_source(source_instance) if source_instance.pending.empty?
          source_instance.pending[key] = value
        end
      end
    end

    # @param trace_query_lazy [nil, Execution::Multiplex]
    def run(trace_query_lazy: nil)
      trace = Fiber[:__graphql_current_multiplex]&.current_trace
      jobs_fiber_limit, total_fiber_limit = calculate_fiber_limit
      job_fibers = []
      next_job_fibers = []
      source_fibers = []
      next_source_fibers = []
      first_pass = true
      manager = spawn_fiber do
        trace&.begin_dataloader(self)
        while first_pass || !job_fibers.empty?
          first_pass = false

          run_pending_steps(trace, job_fibers, next_job_fibers, jobs_fiber_limit, source_fibers, next_source_fibers, total_fiber_limit)

          if !@lazies_at_depth.empty?
            with_trace_query_lazy(trace_query_lazy) do
              if enqueue_next_pending_lazies(@lazies_at_depth)
                job_fibers.unshift(spawn_job_fiber(trace))
                run_pending_steps(trace, job_fibers, next_job_fibers, jobs_fiber_limit, source_fibers, next_source_fibers, total_fiber_limit)
              end
            end
          end
        end

        trace&.end_dataloader(self)
      end

      run_fiber(manager)

      if manager.alive?
        raise "Invariant: Manager fiber didn't terminate properly."
      end

      if !job_fibers.empty?
        raise "Invariant: job fibers should have exited but #{job_fibers.size} remained"
      end
      if !source_fibers.empty?
        raise "Invariant: source fibers should have exited but #{source_fibers.size} remained"
      end

    rescue UncaughtThrowError => e
      throw e.tag, e.value
    end

    def run_fiber(f)
      f.resume
    end

    # @api private
    def lazy_at_depth(depth, lazy)
      @lazies_at_depth[depth] << lazy
    end

    def spawn_fiber
      fiber_vars = get_fiber_variables
      Fiber.new(blocking: !@nonblocking) {
        set_fiber_variables(fiber_vars)
        yield
        cleanup_fiber
      }
    end

    # Pre-warm the Dataloader cache with ActiveRecord objects which were loaded elsewhere.
    # These will be used by {Dataloader::ActiveRecordSource}, {Dataloader::ActiveRecordAssociationSource} and their helper
    # methods, `dataload_record` and `dataload_association`.
    # @param records [Array<ActiveRecord::Base>] Already-loaded records to warm the cache with
    # @param index_by [Symbol] The attribute to use as the cache key. (Should match `find_by:` when using {ActiveRecordSource})
    # @return [void]
    def merge_records(records, index_by: :id)
      records_by_class = Hash.new { |h, k| h[k] = {} }
      records.each do |r|
        records_by_class[r.class][r.public_send(index_by)] = r
      end
      records_by_class.each do |r_class, records|
        with(ActiveRecordSource, r_class).merge(records)
      end
    end

    private

    # Returns true if anything was actually enqueued
    def enqueue_next_pending_lazies(lazies_at_depth)
      smallest_depth = lazies_at_depth.each_key.min
      return false if smallest_depth.nil?

      lazies = lazies_at_depth.delete(smallest_depth)
      return false if lazies.empty?

      lazies.each do |lazy|
        append_job { lazy.value }
      end

      true
    end

    def run_pending_steps(trace, job_fibers, next_job_fibers, jobs_fiber_limit, source_fibers, next_source_fibers, total_fiber_limit)
      while (f = (job_fibers.shift || (((next_job_fibers.size + job_fibers.size) < jobs_fiber_limit) && spawn_job_fiber(trace))))
        if f.alive?
          finished = run_fiber(f)
          if !finished
            next_job_fibers << f
          end
        end
      end
      join_queues(job_fibers, next_job_fibers)

      while (!source_fibers.empty? || !@pending_sources.empty?)
        while (f = source_fibers.shift || (((job_fibers.size + source_fibers.size + next_source_fibers.size + next_job_fibers.size) < total_fiber_limit) && spawn_source_fiber(trace)))
          if f.alive?
            finished = run_fiber(f)
            if !finished
              next_source_fibers << f
            end
          end
        end
        join_queues(source_fibers, next_source_fibers)
      end
    end

    def with_trace_query_lazy(multiplex_or_nil, &block)
      if (multiplex = multiplex_or_nil)
        query = multiplex.queries.length == 1 ? multiplex.queries[0] : nil
        multiplex.current_trace.execute_query_lazy(query: query, multiplex: multiplex, &block)
      else
        yield
      end
    end

    def calculate_fiber_limit
      total_fiber_limit = @fiber_limit || Float::INFINITY
      if total_fiber_limit < 4
        raise ArgumentError, "Dataloader fiber limit is too low (#{total_fiber_limit}), it must be at least 4"
      end
      total_fiber_limit -= 1 # deduct one fiber for `manager`
      # Deduct at least one fiber for sources
      jobs_fiber_limit = total_fiber_limit - 2
      return jobs_fiber_limit, total_fiber_limit
    end

    def join_queues(prev_queue, new_queue)
      @nonblocking && Fiber.scheduler.run
      prev_queue.concat(new_queue)
      new_queue.clear
    end

    def spawn_job_fiber(trace)
      if !@pending_jobs.empty?
        spawn_fiber do
          trace&.dataloader_spawn_execution_fiber(@pending_jobs)
          while job = @pending_jobs.shift
            job.call
          end
          trace&.dataloader_fiber_exit
        end
      end
    end

    def drain_pending_sources
      pending_sources = @pending_sources
      @pending_sources = []
      @pending_source_set.clear

      pending_sources.select!(&:pending?)
      pending_sources.empty? ? nil : pending_sources
    end

    def dequeue_pending_source
      while (source = @pending_sources.shift)
        @pending_source_set.delete(source)
        return source if source.pending?
      end
    end

    def spawn_source_fiber(trace)
      if !@pending_sources.empty?
        spawn_fiber do
          trace&.dataloader_spawn_source_fiber(@pending_sources)
          # This will find sources which were enqueued during `#fetch`:
          while (source = dequeue_pending_source)
            next if !source.pending?
            Fiber[:__graphql_current_dataloader_source] = source
            trace&.begin_dataloader_source(source)
            source.run_pending_keys
            trace&.end_dataloader_source(source)
          end
          trace&.dataloader_fiber_exit
        end
      end
    end
  end
end

require "graphql/dataloader/async_dataloader"
