-- * Metronome IM *
--
-- This file is part of the Metronome XMPP server and is released under the
-- ISC License, please see the LICENSE file in this source package for more
-- information about copyright and licensing.

-- Message Archiving Management for Metronome,
-- This implements XEP-313.

local bare_sessions = metronome.bare_sessions;
local module_host = module.host;

local dt_parse = require "util.datetime".parse;
local st = require "util.stanza";
local jid_bare = require "util.jid".bare;
local jid_join = require "util.jid".join;
local jid_prep = require "util.jid".prep;
local jid_section = require "util.jid".section;
local jid_split = require "util.jid".split;
local collectgarbage, ipairs, now, tonumber, tostring = collectgarbage, ipairs, os.time, tonumber, tostring;

local xmlns = "urn:xmpp:mam:2";
local purge_xmlns = "http://metronome.im/protocol/mam-purge";
local rsm_xmlns = "http://jabber.org/protocol/rsm";
local sid_xmlns = "urn:xmpp:sid:0";

local forbid_purge = module:get_option_boolean("mam_forbid_purge", false);

local mamlib = module:require("mam");
local validate_query = module:require("validate").validate_query;
local initialize_storage, initialize_session_store, initialize_offline_store, save_stores =
	mamlib.initialize_storage, mamlib.initialize_session_store, mamlib.initialize_offline_store, mamlib.save_stores;
local get_prefs, set_prefs = mamlib.get_prefs, mamlib.set_prefs;
local fields_handler, generate_stanzas, process_message, purge_messages =
	mamlib.fields_handler, mamlib.generate_stanzas, mamlib.process_message, mamlib.purge_messages;

local session_stores = mamlib.session_stores;
local offline_stores = mamlib.offline_stores;
local storage = initialize_storage();

-- Handlers

local function remove_session_store(event)
	local user, host = event.session.username, event.session.host;
	local bare_jid = jid_join(user, host);
	local bare_session = bare_sessions[bare_jid];
	if not bare_session then session_stores[bare_jid] = nil; end -- dereference session store.
end

local function save_session_store(event)
	local user, host = event.session.username, event.session.host;
	local bare_jid = jid_join(user, host);
	local user_archive = session_stores[bare_jid];
	if user_archive and user_archive.changed then
		local last_used = user_archive.last_used;
		user_archive.changed, user_archive.last_used = nil, nil;
		storage:set(user, user_archive);
		user_archive.last_used = last_used;
	end
end

local function delete_store(event)
	local user, host = event.username, event.host;
	local bare_jid = jid_join(user, host);
	session_stores[bare_jid] = nil;
	storage:set(user, nil);
end

local function process_inbound_messages(event)
	process_message(event);
end

local function process_outbound_messages(event)
	process_message(event, true);
end

local function feature_handler(event)
	local stanza = event.stanza;
	stanza:tag("feature", { var = xmlns }):up();
	stanza:tag("feature", { var = sid_xmlns }):up();
	stanza:tag("feature", { var = purge_xmlns }):up();
end

local function prefs_handler(event)
	local origin, stanza = event.origin, event.stanza;
	local bare_session = bare_sessions[jid_bare(origin.full_jid)];

	if not bare_session.archiving then
		initialize_session_store(origin.username);
	end

	bare_session.archiving.last_used = now();

	if stanza.attr.type == "get" then
		local reply = st.reply(stanza);
		reply:add_child(get_prefs(bare_session.archiving));
		return origin.send(reply);
	else
		local _prefs = stanza.tags[1];
		local reply = set_prefs(stanza, bare_session.archiving);
		return origin.send(reply);
	end
end

local function purge_handler(event)
	local origin, stanza = event.origin, event.stanza;
	local purge = stanza.tags[1];
	local bare_jid = jid_bare(origin.full_jid);
	
	local bare_session = bare_sessions[bare_jid];
	local archive = bare_session.archiving;

	if not archive then
		archive = initialize_session_store(origin.username);
	end

	archive.last_used = now();
	
	if forbid_purge then
		return origin.send(st.error_reply(stanza, "cancel", "not-allowed", "Purging message archives is not allowed"));
	end
	
	local id, jid, start, fin =
		purge:get_child_text("id"), purge:get_child_text("jid"), purge:get_child_text("start"), purge:get_child_text("end");

	local vjid, vstart, vfin = (jid and jid_prep(jid)), (start and dt_parse(start)), (fin and dt_parse(fin));
	if (start and not vstart) or (fin and not vfin) or (jid and not vjid) then
		return origin.send(st.error_reply(stanza, "modify", "bad-request", "Supplied parameters failed verification"));
	end
	
	purge_messages(archive, id, vjid, vstart, vfin);
	module:log("debug", "%s purged Archives", bare_jid);
	return origin.send(st.reply(stanza));
end

local function query_handler(event)
	local origin, stanza = event.origin, event.stanza;
	local query = stanza.tags[1];
	local qid = query.attr.queryid;
		
	local bare_session = bare_sessions[jid_bare(origin.full_jid)];
	local archive = bare_session.archiving;

	if not archive then
		archive = initialize_session_store(origin.username);
	end

	archive.last_used = now();

	local start, fin, with, after, before, max, index;
	local ok, ret = validate_query(stanza, query, qid);
	if not ok then
		return origin.send(ret);
	else
		start, fin, with, after, before, max, index =
			ret.start, ret.fin, ret.with, ret.after, ret.before, ret.max, ret.index;
	end
	
	local messages, rq, count = generate_stanzas(archive, start, fin, with, max, after, before, index, qid);
	if not messages then -- RSM item-not-found
		module:log("debug", "MAM Query RSM parameters were out of bounds");
		local rsm_error = st.error_reply(stanza, "cancel", "item-not-found");
		rsm_error:add_child(query);
		return origin.send(rsm_error);
	end

	local reply = st.reply(stanza):add_child(rq);

	for _, message in ipairs(messages) do
		message.attr.to = origin.full_jid;
		origin.send(message);
	end
	origin.send(reply);

	module:log("debug", "MAM query %s completed (returned messages: %s)",
		qid and qid or "without id", count == 0 and "none" or tostring(count));
	return true;
end

function module.unload()
	save_stores();
	-- remove all caches from bare_sessions.
	for bare_jid, bare_session in pairs(bare_sessions) do
		local host = jid_section(bare_jid, "host");
		if host == module_host then bare_session.archiving = nil; end
	end
	module:remove_all_timers();
	module:add_timer(5, function() collectgarbage(); end);
end

module:hook("mam-get-store", function(user)
	local user_jid = jid_join(user, module_host);
	local session_store = session_stores[user_jid];
	local offline_store = offline_stores[user_jid];
	if session_store then
		return session_store;
	else
		if bare_sessions[user_jid] then
			return initialize_session_store(user);
		elseif offline_store then
			return offline_store;
		else
			return initialize_offline_store(user, true);
		end
	end
end);

module:hook("pre-resource-unbind", save_session_store, 30);
module:hook("resource-unbind", remove_session_store, 30);
module:hook("user-deleted", delete_store);

module:hook("message/bare", process_inbound_messages, 30);
module:hook("pre-message/bare", process_outbound_messages, 30);
module:hook("message/full", process_inbound_messages, 30);
module:hook("pre-message/full", process_outbound_messages, 30);

module:hook("account-disco-info", feature_handler, 35);
module:hook("iq/self/"..xmlns..":prefs", prefs_handler);
module:hook("iq-set/self/"..purge_xmlns..":purge", purge_handler);
module:hook("iq-set/self/"..xmlns..":query", query_handler);
module:hook("iq-get/self/"..xmlns..":query", fields_handler);

module:hook_global("server-stopping", save_stores);
