-- Copyright (C) 2014-2015 Kim Alvefur

--

-- This file is MIT/X11 licensed.

local tonumber = tonumber;

local tostring = tostring;

local type = type;

local io_open = io.open;

local os_remove = os.remove;

local os_rename = os.rename;

local noop = function () end

local id = function (x) return x end

local lfs = require"lfs";

local st = require"util.stanza";

local xml = require"util.xml";

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

local datetime = require"util.datetime";

local path_join = require"util.paths".join;

local http = require"util.http";

local encode, decode = import("util.serialization", "serialize", "deserialize");

local decode_file = encoder.decode_file;

local log = module._log;

local base_path = prosody.paths.data;

local suffix = ".dat";

-- No path escaping, assume no forbidden chars in names

local function getpath(...)

	return path_join(base_path, module.host, ...);

end

local did_mkdir_already = {};

local function mkdir(path)

	if not did_mkdir_already[path] then

		did_mkdir_already[path] = true;

		return lfs.mkdir(path);

	end

end

function module.load()

	mkdir(getpath()); -- Create base dir for host

end

-- Cleanup

local function cleanup(scratch, f, msg)

	if f then f:close(); end

	os_remove(scratch);

	return nil, msg;

end

local function atomic_store(filename, data)

	local scratch = filename.."~";

	local f, ok, msg;

	f, msg = io_open(scratch, "wb");

	if not f then return cleanup(scratch, f, msg) end

	ok, msg = f:write(data);

	if not ok then return cleanup(scratch, f, msg) end

	ok, msg = f:close();

	if not ok then return cleanup(scratch, f, msg) end

	return os_rename(scratch, filename);

end

local function read_fh(fh, max_length)

	return decode(fh:read(max_length or "*a"));

end

local function read_file(path, max_length)

	local f, err = io_open(path, "rb");

	if not f then

		return f, err;

	end

	local data = read_fh(fh, max_length);

	f:close();

	return data;

end

local function log_error(path, store, err, user)

	local mode = lfs.attributes(path, "mode");

	if not mode then

		log("debug", "Assuming empty %s storage ('%s') for user: %s", store, err, user or "nil");

		return nil;

	else -- file exists, but can't be read

		-- TODO more detailed error checking and logging?

		log("error", "Failed to load %s storage ('%s') for user: %s", store, err, user or "nil");

		return nil, "Error reading storage";

	end

end

local function h(s)

	return sha256(s, true):sub(1, 2);

end

local function new_keyval_store(store)

	local driver = { store = store, typ = "keyval" };

	function driver:get(user)

		user = user or "@";

		local path = getpath(user, store) .. suffix;

		local data, err = read_file(path);

		if not data then return log_error(path, store, err, user or "<nil>") end

		return data;

	end

	function driver:set(user, data)

		user = user or "@";

		mkdir(getpath(user));

		local path = getpath(user, store) .. suffix;

		if data == nil then

			if lfs.attributes(path, "mode") == "file" then

				return os_remove(path);

			end

			return true;

		end

		data = encode(data);

		return atomic_store(path, data);

	end

	function driver:users()

		local next, state = lfs.dir(getpath());

		return function (state)

			for node in next, state do

				if lfs.attributes(getpath(node, store) .. suffix, "mode") == "file" then

					return node;

				end

			end

		end, state;

	end

	return driver;

end

local function new_map_store(store, as_keyval)

	local driver = { store = store, typ = "map" };

	-- FIXME Keys may be encoded to values longer than the file system supports

	function driver:get(user, key)

		assert(key ~= nil, "map key can't be nil");

		user = user or "@";

		key = http.urlencode(encode(key));

		local path = getpath(user, store, key) .. suffix;

		local data, err = read_file(path);

		if not data then return log_error(path, store, err, user or "<nil>") end

		return data;

	end

	function driver:set(user, key, data)

		assert(key ~= nil, "map key can't be nil");

		user = user or "@";

		key = http.urlencode(encode(key));

		mkdir(getpath(user));

		mkdir(getpath(user, store));

		local path = getpath(user, store, key) .. suffix;

		if data == nil then

			return os_remove(path);

		end

		data = encode(data);

		return atomic_store(path, data);

	end

	function driver:keys(user)

		user = user or "@";

		local ok, iter, state, var = pcall(lfs.dir, getpath(user, store));

		if not ok then

			log("warn", iter);

			return noop;

		end

		return function (state, var)

			local item;

			repeat

				item = iter(state, var);

				if item == nil then return end

			until lfs.attributes(getpath(user,store, item), "mode") == "file" and item:sub(-#suffix) == suffix;

			return decode(http.urldecode(item:sub(1, -#suffix-1)));

		end, state, var;

	end

	function driver:pairs(user)

		local iter, state, var = self:keys(user);

		return function (state, var)

			local key = iter(state, var);

			if key == nil then return end

			return key, self:get(user, key);

		end, state, var;

	end

	function driver:empty(user)

		for key in self:keys(user) do

			self:set(user, key, nil);

		end

		return true;

	end

	function driver:collect(user)

		local data = {};

		for k,v in self:pairs(user) do

			data[k] = v;

		end

		return data;

	end

	if as_keyval then

		local keyval_driver = new_keyval_store(store);

		function keyval_driver:get(user)

			return driver:collect(user);

		end

		function keyval_driver:set(user, data)

			driver:empty(user);

			if data ~= nil then

				for k, v in pairs(data) do

					driver:set(user, k, v);

				end

			end

		end

		return keyval_driver;

	end

	return driver;

end

-- date utilities

local function tomorrow(date)

	return datetime.date(datetime.parse(date .. "T12:00:00Z") + 86400);

end

local function yesterday(date)

	return datetime.date(datetime.parse(date .. "T12:00:00Z") - 86400);

end

local function new_archive_store(store)

	local driver = { store = store, typ = "archive" };

	local fallocate = require"util.pposix".fallocate or function (f, o, l)

		local ok, err = f:write(string.rep(" ", l));

		if not ok then return ok, err end

		return f:seek("set", o);

	end

	local index_storage = new_map_store(store);

	local _get = index_storage.get;

	function index_storage:get(user, key)

		local ok, err = _get(self, user, key);

		if not ok then

			log("debug", tostring(err));

		end

		return ok, err;

	end

	local datelists = setmetatable({}, { __mode = 'v'; });

	local function get_datelist(user, with)

		local key = with and (user..'\0'..with) or user;

		local list = datelists[key];

		if list then return list end

		if with then

			local datemap = index_storage:get(user, with);

			if datemap then

				list = {};

				local i = 1;

				for date in pairs(datemap) do

					list[i], i = date, i+1;

				end

			end

		end

		if not list then

			list = {};

			local ok, iter, state, var = pcall(lfs.dir, getpath(user, store));

			if not ok then return list; end

			local i = 1;

			local filename, mode, date;

			for item in iter, state, var do

				filename = getpath(user, store, item);

				mode = lfs.attributes(filename, "mode");

				if mode == "file" and item:match("^%d%d%d%d%-%d%d%-%d%d") and item:sub(-#suffix) == suffix then

					list[i], i = item:sub(1,10), i+1;

				end

			end

		end

		table.sort(list);

		for i = 1, #list do

			list[ list[i] ] = i;

		end

		datelists[key] = list;

		return list;

	end

	function driver:append(user, key, when, with, value)

		user = user or "@";

		local ok;

		local date = datetime.date(when);

		mkdir(getpath(user));

		mkdir(getpath(user, store));

		local filename = getpath(user, store, date) .. suffix;

		local f, msg = io_open(filename, "r+b");

		if not f then

			f, msg = io_open(filename, "wb");

		end

		if not f then

			return f, msg;

		end

		local mt = getmetatable(value);

		local typ = type(value);

		--[[

		if mt.__freeze then

			value, typ = mt.__freeze(value);

		else--]]

		if mt == st.stanza_mt then

			value, typ = st.preserialize(value), "stanza";

			-- if false then

				-- value, typ = tostring(value), "xml";

			-- else

			-- end

		end

		local pos = f:seek("end");

		key = ("%sp%x"):format(date, pos);

		key = key .. "r".. sha256(user..key, true):sub(1, 8);

		local metadata = encode({ key, datetime.datetime(when), with, typ });

		local data = encode(value);

		local meta_len, data_len = #metadata, #data;

		local meta_ls, data_ls = ("%d"):format(meta_len), ("%d"):format(data_len);

		local total_ls = ("%d"):format(meta_len + #meta_ls + data_len + #data_ls + 4); -- two newlines?

		-- ok, msg = fallocate(f, pos, data_len + #len_str * 2 + 3);

		-- if not ok then return ok, msg; end

		ok, msg = f:write(meta_ls, "\n", metadata, "\n", data_ls, "\n", data, "\n", total_ls, "\n");

		if not ok then return ok, msg; end

		ok, msg = f:close();

		if not ok then return ok, msg; end

		datelists[user] = nil;

		if with then datelists[user .. "\0" .. with] = nil end

		local index = index_storage:get(user, with) or {};

		if not index[date] then

			index[date] = true;

			index_storage:set(user, with, index);

		end

		return key;

	end

	local unfreeze = {

		xml = require"util.xml".parse;

		stanza = st.deserialize;

	}

	local function valid_key(user, key)

		if not key then return nil, "no-key"; end

		if sha256(user..key:match("^(.-)r"), 1):sub(1, 8) ~= key:match("r(.*)$") then

			log("warn", "invalid-key-signature");

			-- return nil, "invalid-key-signature";

		end

		local when = datetime.parse(key:sub(1, 10).."T00:00:00Z");

		if not when then return nil, "invalid-key-date"; end

		local pos = key:match("^p(%x+)r", 11);

		pos = tonumber(pos, 16);

		if not pos then return nil, "invalid-key-offset"; end

		return when, pos;

	end

	function driver:get(user, key)

		local date, seek = valid_key(user, key);

		if not date then return nil, seek; end

		date = datetime.date(date);

		local filename = getpath(user, store, date) .. suffix;

		local fh, err = io_open(filename, "rb");

		if not fh then

			return nil, err;

		end

		fh:seek("set", seek);

		local len = fh:read("*l");

		local meta = fh:read(tonumber(len));