require 'json'

Puppet::Functions.create_function(:ldap_data) do
  dispatch :ldap_data do
    param 'Hash', :options
    param 'Puppet::LookupContext', :context
  end

  def ldap_data(options, context)
    begin
      require 'ldap'
      require 'puppet/util/ldap/connection'
    rescue LoadError
      context.not_found
      return
    end

    if !context.cache_has_key("ldap_lookup")
      begin
        conn = Puppet::Util::Ldap::Connection.instance
        conn.start
        connection = conn.connection
      rescue ::LDAP::ResultError => e
        raise Puppet::ParseError, ("ldapquery(): LDAP ResultError - #{e.message}")
      end

      host = Facter.value('ec2_metadata')["hostname"]
      base = Puppet[:ldapbase]
      scope  = ::LDAP::LDAP_SCOPE_SUBTREE
      filter = "(objectclass=*)"

      data = {}
      data_array = connection.search2(base, scope, filter, attrs=["immaePuppetJson", "dn"]).map do |entry|
        [entry["dn"].first, entry["immaePuppetJson"] || []]
      end.sort_by do |dn, json|
        if dn == "ou=roles,ou=hosts,dc=immae,dc=eu"
          [0, dn]
        elsif dn.end_with?("ou=roles,ou=hosts,dc=immae,dc=eu")
          [1, dn]
        else
          [2, dn]
        end
      end.to_h

      data_array.each do |dn, jsons|
        jsons.each do |json|
          data.merge!(JSON.parse(json))
        end
      end

      context.cache("ldap_lookup", data)
    end

    context.cached_value("ldap_lookup")
  end
end