local st = require "util.stanza";

local affils = {
	none = 0;
	member = 1;
	administrator = 2;
	owner = 3;
};

local function get_restrict_pm(room)
	return room._data.restrict_pm or 'none';
end

local function set_restrict_pm(room, affil)
	if get_restrict_pm(room) == affil then return false; end
	room._data.restrict_pm = affil;
	return true;
end

local function can_pm_anyone(room)
	return room._data.restrict_pm_to or false;
end

local function set_pm_anyone(room, val)
	if can_pm_anyone(room) == val then return false; end
	room._data.restrict_pm_to = val;
	return true;
end

local function is_visitor_pm_off(room)
	return room._data.restrict_pm_visitor or false;
end

local function set_visitor_pm_off(room, val)
	if is_visitor_pm_off(room) == val then return false; end
	room._data.restrict_pm_visitor = val;
	return true;
end

module:hook("muc-config-form", function(event)
	local affilpm = get_restrict_pm(event.room);
	table.insert(event.form, {
		name = 'muc#restrict_pm';
		type = 'list-single';
		label = 'Allow PMs from';
		options = {
			{ value = 'owner', label = 'Owner', default = affilpm == 'owner' },
			{ value = 'administrator', label = 'Administrators', default = affilpm == 'administrator' },
			{ value = 'member', label = 'Members', default = affilpm == 'member' },
			{ value = 'none', label = 'Everyone', default = affilpm == 'none' }
		}
	});

	table.insert(event.form, {
		name = 'muc#restrict_pm_to';
		type = 'boolean';
		label = 'Allow PMs to everyone';
		value = can_pm_anyone(event.room);
	});

	table.insert(event.form, {
		name = 'muc#restrict_pm_visitor';
		type = 'boolean';
		label = 'Disable PMs from Visitors';
		value = is_visitor_pm_off(event.room);
	});
end);

module:hook("muc-config-submitted/muc#restrict_pm", function(event)
	if set_restrict_pm(event.room, event.value) then
		event.status_codes["104"] = true;
	end
end);

module:hook("muc-config-submitted/muc#restrict_pm_to", function(event)
	if set_pm_anyone(event.room, event.value) then
		event.status_codes["104"] = true;
	end
end);

module:hook("muc-config-submitted/muc#restrict_pm_visitor", function(event)
	if set_visitor_pm_off(event.room, event.value) then
		event.status_codes["104"] = true;
	end
end);

local function can_user_pm(room, jid)
	local affil, pmval = affils[room:get_affiliation(jid) or 'none'], affils[get_restrict_pm(room)];
	if affil >= pmval then return true; end
	return false;
end

module:hook("muc-private-message", function(event)
	local stanza, room = event.stanza, event.room;
	local from_occupant = room:get_occupant_by_nick(stanza.attr.from);
	local to_occupant = room:get_occupant_by_nick(stanza.attr.to);

	-- To self is always okay
	if to_occupant.bare_jid == from_occupant.bare_jid then return; end

	-- To moderation is okay
	if to_occupant and to_occupant.role == 'moderator' then return; end

	if not is_visitor_pm_off(room) or (from_occupant == nil or from_occupant.role ~= 'visitor') then
		-- If visitors disabled
		if can_user_pm(room, from_occupant.bare_jid) and (can_pm_anyone(room) or can_user_pm(room, to_occupant.bare_jid)) then
			return;
		end;
	end

	room:route_to_occupant(
		from_occupant,
		st.error_reply(stanza, "cancel", "policy-violation", "Private messages are disabled", room.jid)
	);
	return false;
end, 1);