#!/usr/bin/env ruby

#---
# Copyright 2003-2013 by Jim Weirich (jim.weirich@gmail.com).
# All rights reserved.

# Permission is granted for use, copying, modification, distribution,
# and distribution of modified versions of this work as long as the
# above copyright notice is included.
#+++

require 'flexmock/noop'
require 'flexmock/expectation_builder'

class FlexMock

  # PartialMockProxy is used to mate the mock framework to an existing
  # object. The object is "enhanced" with a reference to a mock object
  # (stored in <tt>@flexmock_proxy</tt>). When the +should_receive+
  # method is sent to the proxy, it overrides the existing object's
  # method by creating singleton method that forwards to the mock.
  # When testing is complete, PartialMockProxy will erase the mocking
  # infrastructure from the object being mocked (e.g. remove instance
  # variables and mock singleton methods).
  #
  class PartialMockProxy
    include Ordering

    attr_reader :mock

    # Boxing of the flexmock proxy
    #
    # It is managed as a stack in order to allow to setup containers recursively
    # (as e.g. FlexMock.use( ... ) checks)
    class ProxyBox
      attr_reader :stack

      Element = Struct.new :proxy, :container

      def initialize
        @stack = [Element.new]
      end

      # Tests whether the given container is the one on which the current proxy
      # acts
      def container
        stack.last.container
      end

      def proxy
        stack.last.proxy
      end

      def push(proxy, container)
        stack.push(Element.new(proxy, container))
      end

      def pop
        if !stack.empty?
          stack.pop
        end
      end

      def empty?
        stack.size == 1
      end
    end

    # Make a partial mock proxy and install it on the target +obj+.
    def self.make_proxy_for(obj, container, name, safe_mode)
      name ||= "flexmock(#{obj.class.to_s})"
      if !obj.instance_variable_defined?("@flexmock_proxy")
        proxy_box = obj.instance_variable_set("@flexmock_proxy", ProxyBox.new)
      else
        proxy_box = obj.instance_variable_get("@flexmock_proxy")
      end

      if proxy_box.container != container
        if !proxy_box.empty?
          parent_proxy, _ = proxy_box.proxy
          parent_mock = parent_proxy.mock
        end

        mock  = FlexMock.new(name, container, parent: parent_mock)
        proxy = PartialMockProxy.new(obj, mock, safe_mode, parent: parent_proxy)
        proxy_box.push(proxy, container)
      end
      proxy_box.proxy
    end

    # The following methods are added to partial mocks so that they
    # can act like a mock.

    MOCK_METHODS = [
      :should_receive, :new_instances, :should_expect,
      :should_receive_with_location,
      :flexmock_get,   :flexmock_teardown, :flexmock_verify,
      :flexmock_received?, :flexmock_calls, :flexmock_find_expectation,
      :invoke_original
    ]

    # Initialize a PartialMockProxy object.
    def initialize(obj, mock, safe_mode, parent: nil)
      @obj = obj
      @mock = mock
      @proxy_definition_module = nil
      @parent = parent
      @initialize_override = nil

      unless safe_mode
        add_mock_method(:should_receive)
        MOCK_METHODS.each do |sym|
          unless @obj.respond_to?(sym)
            add_mock_method(sym)
          end
        end
      end
    end

    # Get the mock object for the partial mock.
    def flexmock_get
      @mock
    end

    def push_flexmock_container(container)
      @mock.push_flexmock_container(container)
    end

    def pop_flexmock_container
      @mock.pop_flexmock_container
    end

    # :call-seq:
    #    should_receive(:method_name)
    #    should_receive(:method1, method2, ...)
    #    should_receive(:meth1 => result1, :meth2 => result2, ...)
    #
    # Declare that the partial mock should receive a message with the given
    # name.
    #
    # If more than one method name is given, then the mock object should
    # expect to receive all the listed melthods.  If a hash of method
    # name/value pairs is given, then the each method will return the
    # associated result.  Any expectations applied to the result of
    # +should_receive+ will be applied to all the methods defined in the
    # argument list.
    #
    # An expectation object for the method name is returned as the result of
    # this method.  Further expectation constraints can be added by chaining
    # to the result.
    #
    # See Expectation for a list of declarators that can be used.
    def should_receive(*args, **kw)
      flexmock_define_expectation(caller, *args, **kw)
    end

    def should_expect(*args)
      yield Recorder.new(self)
    end

    # Invoke the original of a mocked method
    #
    # Usually called in a #and_return statement
    def invoke_original(m, *args, **kw, &block)
      flexmock_invoke_original(m, args, kw, block)
    end

    # Whether the given method's original definition has been stored
    def find_original_method(m)
      it =
        if m.respond_to?(:to_str) || m.respond_to?(:to_sym)
          @obj.method(m)
        else
          m
        end

      while it && (it.owner != @proxy_definition_module)
        it = it.super_method
      end

      return unless it
      while it && it.owner.kind_of?(ProxyDefinitionModule)
        it = it.super_method
      end
      it
    rescue NameError => e
      raise unless e.name == m
    end

    # Whether the given method's original definition has been stored
    def original_method(m)
      unless (m = find_original_method(m))
        raise ArgumentError, "no original method for #{m}"
      end
      m
    end

    # Whether the given method's original definition has been stored
    def has_original_method?(m)
      find_original_method(m)
    end

    # Whether the given method is already being proxied
    def has_proxied_method?(m)
      @proxy_definition_module &&
          @proxy_definition_module.method_defined?(m)
    end

    def flexmock_plain_new_method?(m)
      m.name == :new && m.owner == Class
    end

    def flexmock_define_expectation(location, *args, **kw)
      EXP_BUILDER.parse_should_args(self, args, kw) do |method_name|
        if !has_proxied_method?(method_name)
          define_proxy_method(method_name)
        end
        ex = @mock.flexmock_define_expectation(location, method_name)
        if FlexMock.partials_verify_signatures
          if (existing_method = find_original_method(method_name))
            if flexmock_plain_new_method?(existing_method)
              # Look for the signature of `initialize` instead
              ex.with_signature_matching(@obj.instance_method(:initialize))
            else
              ex.with_signature_matching(existing_method)
            end
          end
        end
        ex.mock = self
        ex
      end
    end

    def flexmock_find_expectation(*args, **kw, &block)
      @mock.flexmock_find_expectation(*args, **kw, &block)
    end

    def add_mock_method(method_name)
      proxy_module_eval do
        define_method(method_name) { |*args, **kw, &block|
          proxy = __flexmock_proxy or
            fail "Missing FlexMock proxy " +
                 "(for method_name=#{method_name.inspect}, self=\#{self})"
          proxy.send(method_name, *args, **kw, &block)
        }
      end
    end

    # :call-seq:
    #   new_instances.should_receive(...)
    #   new_instances { |instance|  instance.should_receive(...) }
    #
    # new_instances is a short cut method for overriding the behavior of any
    # new instances created via a mocked class object.
    #
    # By default, new_instances will mock the behaviour of the :new
    # method.  If you wish to mock a different set of class methods,
    # just pass a list of symbols to as arguments.  (previous versions
    # also mocked :allocate by default.  If you need :allocate to be
    # mocked, just request it explicitly).
    #
    # For example, to stub only objects created by :make (and not
    # :new), use:
    #
    #    flexmock(ClassName).new_instances(:make).should_receive(...)
    #
    def new_instances(*allocators, &block)
      fail ArgumentError, "new_instances requires a Class to stub" unless
        Class === @obj
      location = caller
      allocators = [:initialize] if allocators.empty?

      expectation_recorder = ExpectationRecorder.new

      if allocators.delete(:initialize)
        initialize_stub(expectation_recorder, block)
      end

      allocators.each do |allocate_method|
        flexmock_define_expectation(location, allocate_method).and_return { |*args, **kw|
          create_new_mocked_object(
            allocate_method, args, kw, expectation_recorder, block)
        }
      end
      expectation_recorder
    end

    # Stubs the #initialize method on a class
    def initialize_stub(recorder, expectations_block)
      if !@initialize_override
        expectation_blocks    = @initialize_expectation_blocks = Array.new
        expectation_recorders = @initialize_expectation_recorders = Array.new
        @initialize_override = Module.new do
          define_method :initialize do |*args, **kw, &block|
            if self.class.respond_to?(:__flexmock_proxy) && (mock = self.class.__flexmock_proxy)
              container = mock.flexmock_container
              mock = container.flexmock(self)
              expectation_blocks.each do |b|
                b.call(mock)
              end
              expectation_recorders.each do |r|
                r.apply(mock)
              end
            end
            super(*args, **kw, &block)
          end
        end
        override = @initialize_override
        @obj.class_eval { prepend override }
      end
      if expectations_block
        @initialize_expectation_blocks    << expectations_block
      end
      @initialize_expectation_recorders << recorder
    end

    def initialize_stub?
      !!@initialize_override
    end

    def initialize_stub_remove
      if initialize_stub?
        @initialize_expectation_blocks.clear
        @initialize_expectation_recorders.clear
      end
    end

    # Create a new mocked object.
    #
    # The mocked object is created using the following steps:
    # (1) Allocate with the original allocation method (and args)
    # (2) Pass to the block for custom configuration.
    # (3) Apply any recorded expecations
    #
    def create_new_mocked_object(allocate_method, args, kw, recorder, block)
      new_obj = flexmock_invoke_original(allocate_method, args, kw, nil)
      mock = flexmock_container.flexmock(new_obj)
      block.call(mock) unless block.nil?
      recorder.apply(mock)
      new_obj
    end
    private :create_new_mocked_object

    # Invoke the original definition of method on the object supported by
    # the stub.
    def flexmock_invoke_original(method, args, kw, block)
      if (original_method = find_original_method(method))
        original_method.call(*args, **kw, &block)
      else
        @obj.__send__(:method_missing, method, *args, **kw, &block)
      end
    end

    # Verify that the mock has been properly called.  After verification,
    # detach the mocking infrastructure from the existing object.
    def flexmock_verify
      @mock.flexmock_verify
    end

    # Remove all traces of the mocking framework from the existing object.
    def flexmock_teardown
      if ! detached?
        initialize_stub_remove
        proxy_module_eval do
          methods = instance_methods(false).to_a
          methods.each do |m|
            remove_method m
          end
        end
        if @obj.instance_variable_defined?(:@flexmock_proxy) &&
            (box = @obj.instance_variable_get(:@flexmock_proxy))
          box.pop
        end
        @obj = nil
      end
    end

    # Forward to the mock's container.
    def flexmock_container
      @mock.flexmock_container
    end

    # Forward to the mock
    def flexmock_received?(*args, **kw)
      @mock.flexmock_received?(*args, **kw)
    end

    # Forward to the mock
    def flexmock_calls
      @mock.flexmock_calls
    end

    # Set the proxy's mock container.  This set value is ignored
    # because the proxy always uses the container of its mock.
    def flexmock_container=(container)
    end

    # Forward the request for the expectation director to the mock.
    def flexmock_expectations_for(method_name)
      @mock.flexmock_expectations_for(method_name)
    end

    # Forward the based on request.
    def flexmock_based_on(*args)
      @mock.flexmock_based_on(*args)
    end

    private

    # The singleton class of the object.
    def target_singleton_class
      @obj.singleton_class
    end

    # Evaluate a block (or string) in the context of the singleton
    # class of the target partial object.
    def target_class_eval(*args, **kw, &block)
      target_singleton_class.class_eval(*args, **kw, &block)
    end

    class ProxyDefinitionModule < Module
    end

    # Evaluate a block into the module we use to define the proxy methods
    def proxy_module_eval(*args, **kw, &block)
      if !@proxy_definition_module
        obj = @obj
        @proxy_definition_module = m = ProxyDefinitionModule.new do
          define_method("__flexmock_proxy") do
            if box = obj.instance_variable_get(:@flexmock_proxy)
              box.proxy
            end
          end
        end
        target_class_eval { prepend m }
      end
      @proxy_definition_module.class_eval(*args, **kw, &block)
    end

    # Hide the existing method definition with a singleton defintion
    # that proxies to our mock object.  If the current definition is a
    # singleton, we need to record the definition and remove it before
    # creating our own singleton method.  If the current definition is
    # not a singleton, all we need to do is override it with our own
    # singleton.
    def hide_existing_method(method_name)
      define_proxy_method(method_name)
    end

    # Define a proxy method that forwards to our mock object.  The
    # proxy method is defined as a singleton method on the object
    # being mocked.
    def define_proxy_method(method_name)
      if method_name =~ /=$/
        proxy_module_eval do
          define_method(method_name) do |*args, **kw, &block|
            __flexmock_proxy.mock.__send__(method_name, *args, **kw, &block)
          end
        end
      else
        proxy_module_eval <<-EOD
          def #{method_name}(*args, **kw, &block)
            FlexMock.verify_mocking_allowed!
            __flexmock_proxy.mock.#{method_name}(*args, **kw, &block)
          end
        EOD
      end
    end

    # Have we been detached from the existing object?
    def detached?
      @obj.nil?
    end
  end
end