require 'set'

class Module
  # Error generated by +memoize_delegate+ when a method is called on +nil+ and +allow_nil+
  # option is not used.
  class DelegationError < NoMethodError; end

  RUBY_RESERVED_KEYWORDS = %w(alias and BEGIN begin break case class def defined? do
  else elsif END end ensure false for if in module next nil not or redo rescue retry
  return self super then true undef unless until when while yield)
  DELEGATION_RESERVED_KEYWORDS = %w(_ arg args block)
  DELEGATION_RESERVED_METHOD_NAMES = Set.new(
    RUBY_RESERVED_KEYWORDS + DELEGATION_RESERVED_KEYWORDS
  ).freeze

  # Provides a +memoize_delegate+ class method to easily expose contained objects'
  # public methods as your own.
  #
  # ==== Options
  # * <tt>:to</tt> - Specifies the target object
  # * <tt>:prefix</tt> - Prefixes the new method with the target name or a custom prefix
  # * <tt>:allow_nil</tt> - if set to true, prevents a +NoMethodError+ from being raised
  #
  # The macro receives one or more method names (specified as symbols or
  # strings) and the name of the target object via the <tt>:to</tt> option
  # (also a symbol or string).
  #
  # Delegation is particularly useful with Active Record associations:
  #
  #   class Greeter < ActiveRecord::Base
  #     def hello
  #       'hello'
  #     end
  #
  #     def goodbye
  #       'goodbye'
  #     end
  #   end
  #
  #   class Foo < ActiveRecord::Base
  #     belongs_to :greeter
  #     memoize_delegate :hello, to: :greeter
  #   end
  #
  #   Foo.new.hello   # => "hello"
  #   Foo.new.goodbye # => NoMethodError: undefined method `goodbye' for #<Foo:0x1af30c>
  #
  # Multiple memoize_delegates to the same target are allowed:
  #
  #   class Foo < ActiveRecord::Base
  #     belongs_to :greeter
  #     memoize_delegate :hello, :goodbye, to: :greeter
  #   end
  #
  #   Foo.new.goodbye # => "goodbye"
  #
  # Methods can be memoize_delegated to instance variables, class variables, or constants
  # by providing them as a symbols:
  #
  #   class Foo
  #     CONSTANT_ARRAY = [0,1,2,3]
  #     @@class_array  = [4,5,6,7]
  #
  #     def initialize
  #       @instance_array = [8,9,10,11]
  #     end
  #     memoize_delegate :sum, to: :CONSTANT_ARRAY
  #     memoize_delegate :min, to: :@@class_array
  #     memoize_delegate :max, to: :@instance_array
  #   end
  #
  #   Foo.new.sum # => 6
  #   Foo.new.min # => 4
  #   Foo.new.max # => 11
  #
  # It's also possible to memoize_delegate a method to the class by using +:class+:
  #
  #   class Foo
  #     def self.hello
  #       "world"
  #     end
  #
  #     memoize_delegate :hello, to: :class
  #   end
  #
  #   Foo.new.hello # => "world"
  #
  # Delegates can optionally be prefixed using the <tt>:prefix</tt> option. If the value
  # is <tt>true</tt>, the memoize_delegate methods are prefixed with the name of the object being
  # memoize_delegated to.
  #
  #   Person = Struct.new(:name, :address)
  #
  #   class Invoice < Struct.new(:client)
  #     memoize_delegate :name, :address, to: :client, prefix: true
  #   end
  #
  #   john_doe = Person.new('John Doe', 'Vimmersvej 13')
  #   invoice = Invoice.new(john_doe)
  #   invoice.client_name    # => "John Doe"
  #   invoice.client_address # => "Vimmersvej 13"
  #
  # It is also possible to supply a custom prefix.
  #
  #   class Invoice < Struct.new(:client)
  #     memoize_delegate :name, :address, to: :client, prefix: :customer
  #   end
  #
  #   invoice = Invoice.new(john_doe)
  #   invoice.customer_name    # => 'John Doe'
  #   invoice.customer_address # => 'Vimmersvej 13'
  #
  # If the target is +nil+ and does not respond to the memoize_delegated method a
  # +NoMethodError+ is raised, as with any other value. Sometimes, however, it
  # makes sense to be robust to that situation and that is the purpose of the
  # <tt>:allow_nil</tt> option: If the target is not +nil+, or it is and
  # responds to the method, everything works as usual. But if it is +nil+ and
  # does not respond to the memoize_delegated method, +nil+ is returned.
  #
  #   class User < ActiveRecord::Base
  #     has_one :profile
  #     memoize_delegate :age, to: :profile
  #   end
  #
  #   User.new.age # raises NoMethodError: undefined method `age'
  #
  # But if not having a profile yet is fine and should not be an error
  # condition:
  #
  #   class User < ActiveRecord::Base
  #     has_one :profile
  #     memoize_delegate :age, to: :profile, allow_nil: true
  #   end
  #
  #   User.new.age # nil
  #
  # Note that if the target is not +nil+ then the call is attempted regardless of the
  # <tt>:allow_nil</tt> option, and thus an exception is still raised if said object
  # does not respond to the method:
  #
  #   class Foo
  #     def initialize(bar)
  #       @bar = bar
  #     end
  #
  #     memoize_delegate :name, to: :@bar, allow_nil: true
  #   end
  #
  #   Foo.new("Bar").name # raises NoMethodError: undefined method `name'
  #
  # The target method must be public, otherwise it will raise +NoMethodError+.
  #
  def memoize_delegate(*methods, to: nil, prefix: nil, allow_nil: nil)
    unless to
      raise ArgumentError, 'Delegation needs a target. Supply an options hash with a :to key as the last argument (e.g. memoize_delegate :hello, to: :greeter).'
    end

    if prefix == true && to =~ /^[^a-z_]/
      raise ArgumentError, 'Can only automatically set the delegation prefix when delegating to a method.'
    end

    method_prefix = \
      if prefix
        "#{prefix == true ? to : prefix}_"
      else
        ''
      end

    location = caller_locations(1, 1).first
    file, line = location.path, location.lineno

    to = to.to_s
    to = "self.#{to}" if DELEGATION_RESERVED_METHOD_NAMES.include?(to)

    methods.each do |method|
      # Attribute writer methods only accept one argument. Makes sure []=
      # methods still accept two arguments.
      definition = (method =~ /[^\]]=$/) ? 'arg' : '*args, &block'

      # The following generated method calls the target exactly once, storing
      # the returned value in a dummy variable.
      #
      # Reason is twofold: On one hand doing less calls is in general better.
      # On the other hand it could be that the target has side-effects,
      # whereas conceptually, from the user point of view, the delegator should
      # be doing one call.
      if allow_nil
        method_def = [
          "def #{method_prefix}#{method}(#{definition})",
          "_ = #{to}",
          "if !_.nil? || nil.respond_to?(:#{method})",
          "  if instance_variable_get('@_memoize_delegate_#{to}_#{method}')",
          "    instance_variable_get('@_memoize_delegate_#{to}_#{method}')",
          "  else",
          "    instance_variable_set('@_memoize_delegate_#{to}_#{method}', _.#{method}(#{definition}))",
          "  end",
          "end",
        "end"
        ].join ';'
      else
        exception = %(raise DelegationError, "#{self}##{method_prefix}#{method} memoize_delegated to #{to}.#{method}, but #{to} is nil: \#{self.inspect}")

        method_def = [
          "def #{method_prefix}#{method}(#{definition})",
          " _ = #{to}",
          "  if instance_variable_get('@_memoize_delegate_#{to}_#{method}')",
          "    instance_variable_get('@_memoize_delegate_#{to}_#{method}')",
          "  else",
          "    instance_variable_set('@_memoize_delegate_#{to}_#{method}', _.#{method}(#{definition}))",
          "  end",
          "rescue NoMethodError => e",
          "  if _.nil? && e.name == :#{method}",
          "    #{exception}",
          "  else",
          "    raise",
          "  end",
          "end"
        ].join ';'
      end

      module_eval(method_def, file, line)
    end
  end
end