local bit = require "prosody.util.bitcompat";

local trie_methods = {};
local trie_mt = { __index = trie_methods };

local function new_node()
	return {};
end

function trie_methods:set(item, value)
	local node = self.root;
	for i = 1, #item do
		local c = item:byte(i);
		if not node[c] then
			node[c] = new_node();
		end
		node = node[c];
	end
	node.terminal = true;
	node.value = value;
end

local function _remove(node, item, i)
	if i > #item then
		if node.terminal then
			node.terminal = nil;
			node.value = nil;
		end
		if next(node) ~= nil then
			return node;
		end
		return nil;
	end
	local c = item:byte(i);
	local child = node[c];
	local ret;
	if child then
		ret = _remove(child, item, i+1);
		node[c] = ret;
	end
	if ret == nil and next(node) == nil then
		return nil;
	end
	return node;
end

function trie_methods:remove(item)
	return _remove(self.root, item, 1);
end

function trie_methods:get(item, partial)
	local value;
	local node = self.root;
	local len = #item;
	for i = 1, len do
		if partial and node.terminal then
			value = node.value;
		end
		local c = item:byte(i);
		node = node[c];
		if not node then
			return value, i - 1;
		end
	end
	return node.value, len;
end

function trie_methods:add(item)
	return self:set(item, true);
end

function trie_methods:contains(item, partial)
	return self:get(item, partial) ~= nil;
end

function trie_methods:longest_prefix(item)
	return select(2, self:get(item));
end

function trie_methods:add_subnet(item, bits)
	item = item.packed:sub(1, math.ceil(bits/8));
	local existing = self:get(item);
	if not existing then
		existing = { bits };
		return self:set(item, existing);
	end

	-- Simple insertion sort
	for i = 1, #existing do
		local v = existing[i];
		if v == bits then
			return; -- Already in there
		elseif v > bits then
			table.insert(existing, v, i);
			return;
		end
	end
end

function trie_methods:remove_subnet(item, bits)
	item = item.packed:sub(1, math.ceil(bits/8));
	local existing = self:get(item);
	if not existing then
		return;
	end

	-- Simple insertion sort
	for i = 1, #existing do
		local v = existing[i];
		if v == bits then
			table.remove(existing, i);
			break;
		elseif v > bits then
			return; -- Stop search
		end
	end

	if #existing == 0 then
		self:remove(item);
	end
end

function trie_methods:has_ip(item)
	item = item.packed;
	local node = self.root;
	local len = #item;
	for i = 1, len do
		if node.terminal then
			return true;
		end

		local c = item:byte(i);
		local child = node[c];
		if not child then
			for child_byte, child_node in pairs(node) do
				if type(child_byte) == "number" and child_node.terminal then
					local bits = child_node.value;
					for j = #bits, 1, -1 do
						local b = bits[j]-((i-1)*8);
						if b ~= 8 then
							local mask = bit.bnot(2^b-1);
							if bit.band(bit.bxor(c, child_byte), mask) == 0 then
								return true;
							end
						end
					end
				end
			end
			return false;
		end
		node = child;
	end
end

local function new()
	return setmetatable({
		root = new_node();
	}, trie_mt);
end

local function is_trie(o)
	return getmetatable(o) == trie_mt;
end

return {
	new = new;
	is_trie = is_trie;
};