local filters = require "prosody.util.filters";

local jid = require "prosody.util.jid";

local set = require "prosody.util.set";

local client_watchers = {};

-- active_filters[session] = {

--   filter_func = filter_func;

--   downstream = { cb1, cb2, ... };

-- }

local active_filters = {};

local function subscribe_session_stanzas(session, handler, reason)

	if active_filters[session] then

		table.insert(active_filters[session].downstream, handler);

		if reason then

			handler(reason, nil, session);

		end

		return;

	end

	local downstream = { handler };

	active_filters[session] = {

		filter_in = function (stanza)

			module:log("debug", "NOTIFY WATCHER %d", #downstream);

			for i = 1, #downstream do

				downstream[i]("received", stanza, session);

			end

			return stanza;

		end;

		filter_out = function (stanza)

			module:log("debug", "NOTIFY WATCHER %d", #downstream);

			for i = 1, #downstream do

				downstream[i]("sent", stanza, session);

			end

			return stanza;

		end;

		downstream = downstream;

	};

	filters.add_filter(session, "stanzas/in", active_filters[session].filter_in);

	filters.add_filter(session, "stanzas/out", active_filters[session].filter_out);

	if reason then

		handler(reason, nil, session);

	end

end

local function unsubscribe_session_stanzas(session, handler, reason)

	local active_filter = active_filters[session];

	if not active_filter then

		return;

	end

	for i = #active_filter.downstream, 1, -1 do

		if active_filter.downstream[i] == handler then

			table.remove(active_filter.downstream, i);

			if reason then

				handler(reason, nil, session);

			end

		end

	end

	if #active_filter.downstream == 0 then

		filters.remove_filter(session, "stanzas/in", active_filter.filter_in);

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

	end

	active_filters[session] = nil;

end

local function unsubscribe_all_from_session(session, reason)

	local active_filter = active_filters[session];

	if not active_filter then

		return;

	end

	for i = #active_filter.downstream, 1, -1 do

		local handler = table.remove(active_filter.downstream, i);

		if reason then

			handler(reason, nil, session);

		end

	end

	filters.remove_filter(session, "stanzas/in", active_filter.filter_in);

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

	active_filters[session] = nil;

end

local function unsubscribe_handler_from_all(handler, reason)

	for session in pairs(active_filters) do

		unsubscribe_session_stanzas(session, handler, reason);

	end

end

local s2s_watchers = {};

module:hook("s2sin-established", function (event)

	for _, watcher in ipairs(s2s_watchers) do

		if watcher.target_spec == event.session.from_host then

			subscribe_session_stanzas(event.session, watcher.handler, "opened");

		end

	end

end);

module:hook("s2sout-established", function (event)

	for _, watcher in ipairs(s2s_watchers) do

		if watcher.target_spec == event.session.to_host then

			subscribe_session_stanzas(event.session, watcher.handler, "opened");

		end

	end

end);

module:hook("s2s-closed", function (event)

	unsubscribe_all_from_session(event.session, "closed");

end);

local watched_hosts = set.new();

local handler_map = setmetatable({}, { __mode = "kv" });

local function add_stanza_watcher(spec, orig_handler)

	local function filtering_handler(event_type, stanza, session)

		if stanza and spec.filter_spec then

			if spec.filter_spec.with_jid then

				if event_type == "sent" and (not stanza.attr.from or not jid.compare(stanza.attr.from, spec.filter_spec.with_jid)) then

					return;

				elseif event_type == "received" and (not stanza.attr.to or not jid.compare(stanza.attr.to, spec.filter_spec.with_jid)) then

					return;

				end

			end

		end

		return orig_handler(event_type, stanza, session);

	end

	handler_map[orig_handler] = filtering_handler;

	if spec.target_spec.jid then

		local target_is_remote_host = not jid.node(spec.target_spec.jid) and not prosody.hosts[spec.target_spec.jid];

		if target_is_remote_host then

			-- Watch s2s sessions

			table.insert(s2s_watchers, {

				target_spec = spec.target_spec.jid;

				handler = filtering_handler;

				orig_handler = orig_handler;

			});

			-- Scan existing s2sin for matches

			for session in pairs(prosody.incoming_s2s) do

				if spec.target_spec.jid == session.from_host then

					subscribe_session_stanzas(session, filtering_handler, "attached");

				end

			end

			-- Scan existing s2sout for matches

			for local_host, local_session in pairs(prosody.hosts) do --luacheck: ignore 213/local_host

				for remote_host, remote_session in pairs(local_session.s2sout) do

					if spec.target_spec.jid == remote_host then

						subscribe_session_stanzas(remote_session, filtering_handler, "attached");

					end

				end

			end

		else

			table.insert(client_watchers, {

				target_spec = spec.target_spec.jid;

				handler = filtering_handler;

				orig_handler = orig_handler;

			});

			local host = jid.host(spec.target_spec.jid);

			if not watched_hosts:contains(host) and prosody.hosts[host] then

				module:context(host):hook("resource-bind", function (event)

					for _, watcher in ipairs(client_watchers) do

						module:log("debug", "NEW CLIENT: %s vs %s", event.session.full_jid, watcher.target_spec);

						if jid.compare(event.session.full_jid, watcher.target_spec) then

							module:log("debug", "MATCH");

							subscribe_session_stanzas(event.session, watcher.handler, "opened");

						else

							module:log("debug", "NO MATCH");

						end

					end

				end);

				module:context(host):hook("resource-unbind", function (event)

					unsubscribe_all_from_session(event.session, "closed");

				end);

				watched_hosts:add(host);

			end

			for full_jid, session in pairs(prosody.full_sessions) do

				if jid.compare(full_jid, spec.target_spec.jid) then

					subscribe_session_stanzas(session, filtering_handler, "attached");

				end

			end

		end

	else

		error("No recognized target selector");

	end

end

local function remove_stanza_watcher(orig_handler)

	local handler = handler_map[orig_handler];

	unsubscribe_handler_from_all(handler, "detached");

	handler_map[orig_handler] = nil;

	for i = #client_watchers, 1, -1 do

		if client_watchers[i].orig_handler == orig_handler then

			table.remove(client_watchers, i);

		end

	end

	for i = #s2s_watchers, 1, -1 do

		if s2s_watchers[i].orig_handler == orig_handler then

			table.remove(s2s_watchers, i);

		end

	end

end

local function cleanup(reason)

	client_watchers = {};

	s2s_watchers = {};

	for session in pairs(active_filters) do

		unsubscribe_all_from_session(session, reason or "cancelled");

	end

end

return {

	add = add_stanza_watcher;

	remove = remove_stanza_watcher;

	cleanup = cleanup;

};