require 'multi_json'
require 'zk'

module ZooMQ
  class Zookeeper

    attr_reader :globals
    attr_reader :locals

    def initialize(uri)
      @zk = ZK.new(uri)
      @zk.mkdir_p("/alive")
      @globals = KeySpace.new(:global, @zk)
    end

    def announce(identity)
      @locals = KeySpace.new(identity, @zk)
      path = "/alive/#{identity}"
      @zk.block_until_node_deleted(path) if @zk.exists?(path)
      @zk.create(path, mode: :ephemeral)
    end

    def close
      @zk.close!
    end

    def watch(&block)
      @zk.register('/alive', &block)
      servers
    end

    def servers
      @zk.children('/alive', watch: true)
    end

    class KeySpace

      def initialize(type, zk)
        @type, @zk = type, zk
      end

      def path_for(key)
        "/#{@type}/#{key}"
      end

      def get(key)
        value = @zk.get(path_for(key)) rescue nil
        MultiJson.load(value.first) rescue nil
      end

      alias :[] :get

      def set(key, value)
        @zk.create("/#{@type}", ignore: :node_exists)
        value = MultiJson.dump(value)
        if !@zk.create(path_for(key), value, ignore: :node_exists)
          @zk.set(path_for(key), value)
        end
      end

      alias :[]= :set

    end

  end
end