class AASM::StateMachineStore

def clone
  StateMachineStore.new.tap do |store|
    @machines.each_pair do |name, machine|
      store.register(name, machine.clone)
    end
  end
end

def fetch: (Class klass, ?true fallback) -> AASM::StateMachineStore

def fetch(klass, fallback = nil)
  stores[klass.to_s] || fallback && begin
    match = klass.ancestors.find do |ancestor|
      ancestor.include? AASM and stores[ancestor.to_s]
    end
    stores[match.to_s]
  end
end

def initialize
  @machines = Concurrent::Map.new
end

def machine(name)
  @machines[name.to_s]
end

def machine_names
  @machines.keys
end

def register(klass, overwrite = false, state_machine = nil)
  raise "Cannot register #{klass}" unless klass.is_a?(Class)
  case name = template = overwrite
    when FalseClass then stores[klass.to_s] ||= new
    when TrueClass then stores[klass.to_s] = new
    when Class then stores[klass.to_s] = stores[template.to_s].clone
    when Symbol then stores[klass.to_s].register(name, state_machine)
    when String then stores[klass.to_s].register(name, state_machine)
    else raise "Don't know what to do with #{overwrite}"
  end
end

def register(name, machine, force = false)
  raise "Cannot use #{name.inspect} for machine name" unless name.is_a?(Symbol) or name.is_a?(String)
  raise "Cannot use #{machine.inspect} as a machine" unless machine.is_a?(AASM::StateMachine)
  if force
    @machines[name.to_s] = machine
  else
    @machines[name.to_s] ||= machine
  end
end

def stores: () -> Concurrent::Map

def stores
  @stores
end

def unregister(klass)
  stores.delete(klass.to_s)
end