-- Copyright (C) 2017 Thilo Molitor

local max_push_errors = module:get_option_number("push_max_errors", 50);

local host_sessions = prosody.hosts[module.host].sessions;

local push_errors = {};

-- For keeping state across reloads while caching reads

local push_store = (function()

	local store = module:open_store();

	local push_services = {};

	local api = {};

	function api:get(user)

		if not push_services[user] then

			local err;

			push_services[user], err = store:get(user);

			if not push_services[user] and err then

				module:log("warn", "Error reading push notification storage for user '%s': %s", user, tostring(err));

				push_services[user] = {};

				return push_services[user], false;

			end

		end

		if not push_services[user] then push_services[user] = {} end

		return push_services[user], true;

	end

	function api:set(user, data)

		push_services[user] = data;

		local ok, err = store:set(user, push_services[user]);

		if not ok then

			module:log("error", "Error writing push notification storage for user '%s': %s", user, tostring(err));

			return false;

		end

		return true;

	end

	function api:set_identifier(user, push_identifier, data)

		local services = self:get(user);

		services[push_identifier] = data;

		return self:set(user, services);

	end

	return api;

end)();

local function handle_push_error(event)

	local stanza = event.stanza;

	local error_type, condition = stanza:get_error();

	local push_identifier = stanza.attr.id;

	local node = jid.split(stanza.attr.to);

	local from = stanza.attr.from;

	local user_push_services = push_store:get(node);

	if user_push_services[push_identifier] and user_push_services[push_identifier].jid == from and error_type ~= "wait" then

		push_errors[push_identifier] = push_errors[push_identifier] + 1;

		module:log("info", "Got error of type '%s' (%s) for identifier '%s':"

			.."error count for this identifier is now at %s", error_type, condition, push_identifier,

			tostring(push_errors[push_identifier]));

		if push_errors[push_identifier] >= max_push_errors then

			module:log("warn", "Disabling push notifications for identifier '%s'", push_identifier);

			-- remove push settings from sessions

			for _, session in pairs(host_sessions[node].sessions) do

				if session.push_identifier == push_identifier then

					session.push_identifier = nil;

					session.push_settings = nil;

				end

			end

			-- save changed global config

			push_store:set_identifier(node, push_identifier, nil);

			push_errors[push_identifier] = nil;

			-- unhook iq handlers for this identifier

			module:unhook("iq-error/bare/"..push_identifier, handle_push_error);

			module:unhook("iq-result/bare/"..push_identifier, handle_push_success);

		end

	end

	return true;

end

local function handle_push_success(event)

	local stanza = event.stanza;

	local push_identifier = stanza.attr.id;

	local node = jid.split(stanza.attr.to);

	local from = stanza.attr.from;

	local user_push_services = push_store:get(node);

	if user_push_services[push_identifier] and user_push_services[push_identifier].jid == from and push_errors[push_identifier] then

		push_errors[push_identifier] = 0;

		module:log("debug", "Push succeeded, error count for identifier '%s' is now at %s", push_identifier, tostring(push_errors[push_identifier]));

	end

	return true;

end

local function account_dico_info(event)

end

module:hook("account-disco-info", account_dico_info);

local function push_enable(event)

	local push_identifier = push_jid .. "<" .. (push_node or "");

	local push_service = {

	local ok = push_store:set_identifier(origin.username, push_identifier, push_service);

		origin.push_identifier = push_identifier;

		origin.push_settings = push_service;

		origin.log("info", "Push notifications enabled (%s)", tostring(origin.push_identifier));

end

module:hook("iq-set/self/"..xmlns_push..":enable", push_enable);

local function push_disable(event)

	local user_push_services = push_store:get(origin.username);

			origin.log("info", "Push notifications disabled (%s)", tostring(key));

			if origin.push_identifier == key then

				origin.push_identifier = nil;

				origin.push_settings = nil;

			end

			push_errors[key] = nil;

			module:unhook("iq-error/bare/"..key, handle_push_error);

			module:unhook("iq-result/bare/"..key, handle_push_success);

	local ok = push_store:set(origin.username, user_push_services);

	if not ok then

		origin.send(st.error_reply(stanza, "wait", "internal-server-error"));

	else

		origin.send(st.reply(stanza));

	end

end

module:hook("iq-set/self/"..xmlns_push..":disable", push_disable);

local function handle_notify_request(stanza, node, user_push_services)

	if not user_push_services or not #user_push_services then return end

	if stanza and stanza._notify then

		module:log("debug", "Already sent push notification to %s@%s for this stanza, not doing it again", node, module.host);

		return;

	end

	if stanza then

		stanza._notify = true;

	end

	for push_identifier, push_info in pairs(user_push_services) do

		-- increment count and save it

		push_store:set_identifier(node, push_identifier, push_info);

		-- construct push stanza

		local push_publish = st.iq({ to = push_info.jid, from = node .. "@" .. module.host, type = "set", id = push_identifier })

				:tag("publish", { node = push_info.node })

		if stanza and include_sender then

		if stanza and include_body then

		-- send out push

		module:log("debug", "Sending push notification for %s@%s to %s (%s)", node, module.host, push_info.jid, tostring(push_info.node));

		-- handle push errors for this node

		if push_errors[push_identifier] == nil then

			push_errors[push_identifier] = 0;

			module:hook("iq-error/bare/"..push_identifier, handle_push_error);

			module:hook("iq-result/bare/"..push_identifier, handle_push_success);

		end

end

-- small helper function to extract relevant push settings

local function get_push_settings(stanza, session)

	local to = stanza.attr.to;

	local node = to and jid.split(to) or session.username;

	local user_push_services = push_store:get(node);

	return node, user_push_services;

	local node, user_push_services = get_push_settings(event.stanza, event.origin);

	return handle_notify_request(event.stanza, node, user_push_services);

local function process_smacks_stanza(stanza, session)

	if session.push_identifier then

		session.log("debug", "Invoking cloud handle_notify_request for smacks queued stanza...");

		local user_push_services = {[session.push_identifier] = session.push_settings};

		local node = get_push_settings(stanza, session);

		handle_notify_request(stanza, node, user_push_services);

		process_smacks_stanza(queue[i], session);

	filters.add_filter(session, "stanzas/out", process_smacks_stanza);

	local session = event.resumed;

	if session then		-- older smacks module versions send only the "intermediate" session in event.session and no session.resumed one

		filters.remove_filter(session, "stanzas/out", process_smacks_stanza);

		-- this means the counter of outstanding push messages can be reset as well

		if session.push_settings then

			session.push_settings.count = 0;

			push_store:set_identifier(session.username, session.push_identifier, session.push_settings);

		end

	end

	-- process unacked stanzas (handle_notify_request() will only send push requests for new stanzas)

		process_smacks_stanza(queue[i], session);

	end

end

-- archive message added

local function archive_message_added(event)

	-- event is: { origin = origin, stanza = stanza, for_user = store_user, id = id }

	-- only notify for new mam messages when at least one device is only

	if not event.for_user or not host_sessions[event.for_user] then return; end

	local stanza = event.stanza;

	local user_session = host_sessions[event.for_user].sessions;

	local to = stanza.attr.to;

	to = to and jid.split(to) or event.origin.username;

	-- only notify if the stanza destination is the mam user we store it for

	if event.for_user == to then

		local user_push_services = push_store:get(to);

		if not #user_push_services then return end

		-- only notify nodes with no active sessions (smacks is counted as active and handled separate)

		local notify_push_sevices = {};

		for identifier, push_info in pairs(user_push_services) do

			local identifier_found = nil;

			for _, session in pairs(user_session) do

				-- module:log("debug", "searching for '%s': identifier '%s' for session %s", tostring(identifier), tostring(session.push_identifier), tostring(session.full_jid));

				if session.push_identifier == identifier then

					identifier_found = session;

					break;

				end

			end

			if identifier_found then

				identifier_found.log("debug", "Not notifying '%s' of new MAM stanza (session still alive)", identifier);

			else

				notify_push_sevices[identifier] = push_info;

			end

		end

		return handle_notify_request(event.stanza, to, notify_push_sevices);

module:hook("archive-message-added", archive_message_added);

local function send_ping(event)

	local user = event.user;

	local user_push_services = push_store:get(user);

	local push_services = event.push_services or user_push_services;

	return handle_notify_request(nil, user, push_services);

end

-- can be used by other modules to ping one or more (or all) push endpoints

module:hook("cloud-notify-ping", send_ping);

-- TODO: this has to be done on first connect not on offline broadcast, else the counter will be incorrect

-- TODO: it seems this is already done, so this could be safely removed, couldn't it?

-- module:hook("message/offline/broadcast", function(event)

-- 	local origin = event.origin;

-- 	local user_push_services = push_store:get(origin.username);

-- 	if not #user_push_services then return end

-- 

-- 	for _, push_info in pairs(user_push_services) do

-- 		if push_info then

-- 			push_info.count = 0;

-- 		end

-- 	end