-- mod_storage_xmlarchive

-- Copyright (C) 2015-2021 Kim Alvefur

--

-- This file is MIT/X11 licensed.

--

-- luacheck: ignore unused self

local lfs = require "lfs";

local dm = require "core.storagemanager".olddm;

local hmac_sha256 = require"util.hashes".hmac_sha256;

local st = require"util.stanza";

local dt = require"util.datetime";

local new_stream = require "util.xmppstream".new;

local xml = require "util.xml";

local async = require "util.async";

local it = require "util.iterators";

local empty = {};

if not dm.append_raw then

	module:require"datamanager_append_raw";

end

local archive = {};

local archive_mt = { __index = archive };

archive.caps = {

	total = false,

	quota = nil,

	full_id_range = true;

	ids = true;

};

local is_stanza = st.is_stanza or function (s)

	return getmetatable(s) == st.stanza_mt;

end

function archive:append(username, id, data, when, with)

	if not is_stanza(data) then

		module:log("error", "Attempt to store non-stanza object, traceback: %s", debug.traceback());

		return nil, "unsupported-datatype";

	end

	username = username or "@";

	data = tostring(data) .. "\n";

	local day = dt.date(when);

	local ok, err = dm.append_raw(username.."@"..day, self.host, self.store, "xml", data);

	if not ok then

		-- append_raw, unlike list_append, does not log anything on failure atm, so we do so here

		module:log("error", "Unable to write to %s storage ('%s') for user: %s@%s",

			self.store, err, username, module.host)

		return nil, err;

	end

	-- If the day-file is missing then we need to add it to the list of days

	local first_of_day = not lfs.attributes(dm.getpath(username .. "@" .. day, self.host, self.store, "list"));

	local offset = ok and err or 0;

	id = id or day .. "-" .. hmac_sha256(username.."@"..day.."+"..offset, data, true):sub(-16);

	ok, err = dm.list_append(username.."@"..day, self.host, self.store,

		{ id = id, when = dt.datetime(when), with = with, offset = offset, length = #data });

	if ok and first_of_day then

		ok, err = dm.list_append(username, self.host, self.store, day);

	end

	if not ok then

		return nil, err;

	end

	return id;

end

function archive:get(username, id)

	local dates = self:dates(username) or empty;

	local day_idx, item_idx, items = self:_get_idx(username, id, dates);

	if not day_idx then

		return nil, "item-not-found";

	end

	module:log("debug", ":_get_idx(%q, %q) --> %q, %q", username, id, dates[day_idx], items[item_idx]);

	local day = dates[day_idx];

	local item = items[item_idx];

	module:log("debug", "item = %q", item);

	local filename = dm.getpath(username.."@"..day, self.host, self.store, "xml");

	local xmlfile, ferr = io.open(filename, "r");

	if not xmlfile then return nil, ferr; end

	local p,err = xmlfile:seek("set", item.offset);

	if p ~= item.offset or err ~= nil then return nil, err; end

	local data = xmlfile:read(item.length);

	local parsed, perr = xml.parse(data);

	if not parsed then return nil, perr; end

	return parsed, tonumber(item.when) or dt.parse(item.when), item.with;

end

local overwrite = module:get_option("xmlarchive_overwrite", false);

function archive:set(username, id, data, new_when, new_with)

	if not is_stanza(data) then

		module:log("error", "Attempt to store non-stanza object, traceback: %s", debug.traceback());

		return nil, "unsupported-datatype";

	end

	username = username or "@";

	data = tostring(data) .. "\n";

	local dates = self:dates(username) or empty;

	local day_idx, item_idx, items = self:_get_idx(username, id, dates);

	if not day_idx then

		return nil, "item-not-found";

	end

	local day = dates[day_idx];

	local item = items[item_idx];

	local filename = dm.getpath(username.."@"..day, self.host, self.store, "xml");

	local replaced, err = dm.append_raw(username.."@"..day, self.host, self.store, "xml", data);

	if not replaced then return nil, err; end

	local new_offset = err;

	-- default yes or no?

	if overwrite then

		local xmlfile, ferr = io.open(filename, "r+");

		if not xmlfile then return nil, ferr; end

		local p,err = xmlfile:seek("set", item.offset);

		if p ~= item.offset or err ~= nil then return nil, err; end

		local _,err = xmlfile:write((" "):rep(item.length));

		if err ~= nil then return nil, err; end

		local _,err = xmlfile:close();

		if err ~= nil then return nil, err; end

	end

	items[item_idx] = {

		id = id,

		when = new_when or item.when,

		with = new_with or item.with,

		offset = new_offset,

		length = #data,

		replaces = item,

	};

	local ok, err = dm.list_store(username.."@"..day, self.host, self.store, items);

	return ok, err;

end

local function get_nexttick()

	-- COMPAT util.async under Lua 5.1 runs into problems trying to yield across

	-- pcall barriers. Workaround for #1646

	if _VERSION ~= "Lua 5.1" and async.ready() then

		return function ()

			-- slow down

			local wait, done = async.waiter();

			module:add_timer(0, done);

			wait();

		end

	else

		-- no async, no-op

		return function () end

	end

end

function archive:_get_idx(username, id, dates)

	module:log("debug", "Looking for item with id %q", id);

	dates = dates or self:dates(username) or empty;

	local date = id:match("^%d%d%d%d%-%d%d%-%d%d");

	local tick = get_nexttick();

	for d = 1, #dates do

		if not date or date == dates[d] then

			module:log("debug", "Loading index for %s", dates[d]);

			local items = dm.list_load(username .. "@" .. dates[d], self.host, self.store) or empty;

			for i = 1, #items do

				if items[i].id == id then

					module:log("debug", "Found item!");

					return d, i, items;

				end

			end

			if date then

				return; -- Assuming no duplicates

			end

		elseif date and date < dates[d] then

			module:log("debug", "Skipping remaining dates after %s", date);

			return; -- List is assumed to be sorted

		end

		-- insert pauses to allow other processing

		if d % 14 == 0 then tick(); end

	end

	module:log("debug", "Item not found");

end

function archive:find(username, query)

	username = username or "@";

	query = query or empty;

	local result;

	local function cb(_, stanza)

		assert(not result, "Multiple items in chunk");

		result = stanza;

	end

	local stream_session = { notopen = true };

	local stream_callbacks = { handlestanza = cb, stream_ns = "jabber:client", default_ns = "jabber:client" };

	local stream = new_stream(stream_session, stream_callbacks);

	local dates = self:dates(username) or empty;

	local function reset_stream()

		stream:reset();

		stream_session.notopen = true;

		stream:feed(st.stanza("stream", { xmlns = "jabber:client" }):top_tag());

		stream_session.notopen = nil;

	end

	reset_stream();

	local filtered_ids = false;

	local filtered_dates = false;

	if query.ids then

		filtered_ids = {};

		filtered_dates = {};

		for _, id in ipairs(query.ids) do

			filtered_ids[id] = true;

			if filtered_dates then

				local date = id:match("^%d%d%d%d%-%d%d%-%d%d");

				if date then

					filtered_dates[date] = true;

				else

					-- if any id diverges from the standard then the item could be from any date

					filtered_dates = nil;

				end

			end

		end

	end

	if filtered_dates then

		for i = #dates, 1, -1 do

			if not filtered_dates[ dates[i] ] then

				table.remove(dates, i);

			end

		end

	end

	local limit = query.limit;

	local start_day, step, last_day = 1, 1, #dates;

	local count = 0;

	local rev = query.reverse;

	if query.start then

		local d = dt.date(query.start);

		for i = start_day, last_day, step do

			if dates[i] < d then

				start_day = i + 1;

			elseif dates[i] >= d then

				start_day = i; break;

			end

		end

	end

	if query["end"] then

		local d = dt.date(query["end"]);

		for i = last_day, start_day, -step do

			if dates[i] > d then

				last_day = i - 1;

			elseif dates[i] <= d then

				last_day = i; break;

			end

		end

	end

	local items;

	local first_item, last_item;

	local stop_at_id;

	if rev then

		start_day, step, last_day = last_day, -step, start_day;

		if query.before then

			local before_day, before_item, items_ = self:_get_idx(username, query.before, dates);

			if not before_day then

				return nil, "item-not-found";

			elseif before_day <= start_day then

				if before_item then

					first_item = before_item - 1;

				else

					first_item = #items_;

				end

				last_item = 1;

				start_day = before_day;

				items = items_;

			end

		end

		if query.after then

			stop_at_id = query.after;

		end

	elseif query.after then

		local after_day, after_item, items_ = self:_get_idx(username, query.after, dates);

		if not after_day then

			return nil, "item-not-found";

		elseif after_day >= start_day then

			if after_item then

				first_item = after_item + 1;

			else

				first_item = 1;

			end

			last_item = #items_;

			start_day = after_day;

			items = items_;

		end

		if query.before then

			stop_at_id = query.before;

		end

	elseif query.before then

		stop_at_id = query.before;

	end

	local date_open, xmlfile;

	local function read_xml(date, offset, length)

		if xmlfile and date ~= date_open then

			module:log("debug", "Closing XML file for %s", date_open);

			xmlfile:close();

			xmlfile = nil;

		end

		if not xmlfile then

			date_open = date;

			local filename = dm.getpath(username .. "@" .. date, self.host, self.store, "xml");

			local ferr;

			xmlfile, ferr = io.open(filename);

			if not xmlfile then

				module:log("error", "Error: %s", ferr);

				return nil, ferr;

			end

			module:log("debug", "Opened XML file %s", filename);

		end

		local pos, err = xmlfile:seek("set", offset);

		if pos ~= offset then

			return nil, err or "seek-failed";

		end

		return xmlfile:read(length);

	end

	local tick = get_nexttick();

	return function ()

		if limit and count >= limit then if xmlfile then xmlfile:close() end return; end

		for d = start_day, last_day, step do

			local date = dates[d];

			if not items then

				module:log("debug", "Loading index for %s", date);

				start_day = d;

				items = dm.list_load(username .. "@" .. date, self.host, self.store) or empty;

				if not rev then

					first_item, last_item = 1, #items;

				else

					first_item, last_item = #items, 1;

				end

			end

			local q_with, q_start, q_end = query.with, query.start, query["end"];

			for i = first_item, last_item, step do

				local item = items[i];

				if not item then

					module:log("warn", "data[%q][%d] is nil", date, i);

					break;

				end

				local i_when, i_with = item.when, item.with;

				if type(i_when) == "string" then

					i_when = dt.parse(i_when);