module:set_global();

local adns = require "net.adns";

local map_config = module:get_option("srvinjection") or {};
local map = module:shared "s2s_map"

for host, mapping in pairs(map_config) do
	if type(mapping) == "table" and type(mapping[1]) == "string" and (type(mapping[2]) == "number") then
		local connecthost, connectport = mapping[1], mapping[2] or 5269;
		map[host] = {{
			srv = {
				target = connecthost..".";
				port = connectport;
				priority = 1;
				weight = 0;
			};
		}};
	else
		module:log("warn", "Ignoring invalid SRV injection for host '%s'", host);
		map[host] = nil;
	end
end

local original_lookup = adns.lookup;
function adns.lookup(handler, qname, qtype, qclass)
	if qtype == "SRV" then
		local host = qname:match("^_xmpp%-server%._tcp%.(.*)%.$");
		local mapping = map[host] or map["*"];
		if mapping then
			handler(mapping);
			return;
		end
	elseif qtype == "A" then
		if (qname == "localhost." or qname == "127.0.0.1.") then
			handler({{ a = "127.0.0.1" }});
			return;
		end
		local ip = qname:match("^(%d+.%d+.%d+.%d+).$");
		if ip then
			handler({{ a = ip }});
			return;
		end
	end
	return original_lookup(handler, qname, qtype, qclass);
end

function module.unload()
	adns.lookup = original_lookup;
end