class RSpec::Mocks::AnyInstance::MessageChains

def [](method_name)
  @chains_by_method_name[method_name]
end

def add(method_name, chain)
  @chains_by_method_name[method_name] << chain
  chain
end

def all_expectations_fulfilled?
  @chains_by_method_name.all? do |method_name, chains|
    chains.all? { |chain| chain.expectation_fulfilled? }
  end
end

def each_unfulfilled_expectation_matching(method_name, *args)
  @chains_by_method_name[method_name].each do |chain|
    if !chain.expectation_fulfilled? && chain.matches_args?(*args)
      yield chain
    end
  end
end

def has_expectation?(method_name)
  @chains_by_method_name[method_name].find do |chain|
    ExpectationChain === chain
  end
end

def initialize
  @chains_by_method_name = Hash.new { |h, k| h[k] = [] }
end

def playback!(instance, method_name)
  raise_if_second_instance_to_receive_message(instance)
  @chains_by_method_name[method_name].each do |chain|
    chain.playback!(instance)
  end
end

def raise_if_second_instance_to_receive_message(instance)
  @instance_with_expectation ||= instance if ExpectationChain === instance
  if ExpectationChain === instance && !@instance_with_expectation.equal?(instance)
    raise RSpec::Mocks::MockExpectationError, "Exactly one instance should have received the following message(s) but didn't: #{unfulfilled_expectations.sort.join(', ')}"
  end
end

def received_expected_message!(method_name)
  @chains_by_method_name[method_name].each do |chain|
    chain.expectation_fulfilled!
  end
end

def remove_stub_chains_for!(method_name)
  @chains_by_method_name[method_name].reject! do |chain|
    StubChain === chain
  end
end

def unfulfilled_expectations
  @chains_by_method_name.map do |method_name, chains|
    method_name.to_s if ExpectationChain === chains.last unless chains.last.expectation_fulfilled?
  end.compact
end