#include "lua_filesystem.h"
#include <filesystem>
#include <chrono>

int luaopen_filesystem(lua_State* L)
{
	luaL_newmetatable(L, "fs.file_handle");

	// Push iterator state metatables
	luaL_newmetatable(L, "fs.stdlibcpp.filesystem.directory_iterator");
	lua_pushcfunction(L, lua_fs_foreach_dir_dtor);
	lua_setfield(L, -2, "__gc");

	luaL_newmetatable(L, "fs.stdlibcpp.filesystem.recursive_directory_iterator");
	lua_pushcfunction(L, lua_fs_forall_dir_dtor);
	lua_setfield(L, -2, "__gc");

	
	// Push functions
	lua_newtable(L);

	lua_pushcfunction(L, lua_fs_is_newer);
	lua_setfield(L, -2, "is_newer");

	lua_pushcfunction(L, lua_fs_filter_newer);
	lua_setfield(L, -2, "filter_newer");

	lua_pushcfunction(L, lua_fs_absolute);
	lua_setfield(L, -2, "absolute");

	lua_pushcfunction(L, lua_fs_create_dir);
	lua_setfield(L, -2, "create_dir");

	lua_pushcfunction(L, lua_fs_exists);
	lua_setfield(L, -2, "exists");

	lua_pushcfunction(L, lua_fs_file_directory);
	lua_setfield(L, -2, "dir");

	lua_pushcfunction(L, lua_fs_file_extension);
	lua_setfield(L, -2, "extension");

	lua_pushcfunction(L, lua_fs_file_name);
	lua_setfield(L, -2, "file_name");

	lua_pushcfunction(L, lua_fs_forall_dir);
	lua_setfield(L, -2, "forall");

	lua_pushcfunction(L, lua_fs_foreach_dir);
	lua_setfield(L, -2, "foreach");

	lua_pushcfunction(L, lua_fs_is_dir);
	lua_setfield(L, -2, "is_dir");

	lua_pushcfunction(L, lua_fs_is_file);
	lua_setfield(L, -2, "is_file");

	lua_pushcfunction(L, lua_fs_last_modified);
	lua_setfield(L, -2, "last_modified");

	lua_pushcfunction(L, lua_fs_list_dir);
	lua_setfield(L, -2, "list_dir");

	lua_pushcfunction(L, lua_fs_list_alldirs);
	lua_setfield(L, -2, "list_alldirs");

	lua_pushcfunction(L, lua_fs_relative);
	lua_setfield(L, -2, "relative_to");
	
	lua_pushvalue(L, -1);
	lua_setfield(L, LUA_GLOBALSINDEX, "fs");

	return 1;
}

int lua_fs_list_dir(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a path to make the second argument relative to");
	const char* fpath = lua_tostring(L, 1);
	lua_newtable(L);
	uint32_t i = 1;
	for (const std::filesystem::path& p : std::filesystem::directory_iterator(fpath))
	{
		const std::string x = p.string();
		lua_pushlstring(L, x.c_str(), x.size());
		lua_rawseti(L, -2, i);
	}
	return 1;
}

int lua_fs_list_alldirs(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a path to make the second argument relative to");
	const char* fpath = lua_tostring(L, 1);
	lua_newtable(L);
	uint32_t i = 1;
	for (const std::filesystem::path& p : std::filesystem::recursive_directory_iterator(fpath))
	{
		const std::string x = p.string();
		lua_pushlstring(L, x.c_str(), x.size());
		lua_rawseti(L, -2, i);
		i++;
	}
	return 1;
}

int lua_fs_foreach_dir(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a path");
	const char* path = lua_tostring(L, 1);

	void* ud = lua_newuserdata(L, sizeof(std::filesystem::directory_iterator));
	new (ud) std::filesystem::directory_iterator(path);

	luaL_getmetatable(L, "fs.stdlibcpp.filesystem.directory_iterator");
	lua_setmetatable(L, -2);
	lua_pushcclosure(L, lua_fs_foreach_dir_next, 1);
	return 1;
}

int lua_fs_forall_dir(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a path");
	const char* path = lua_tostring(L, 1);

	void* ud = lua_newuserdata(L, sizeof(std::filesystem::recursive_directory_iterator));
	new (ud) std::filesystem::recursive_directory_iterator(path);

	luaL_getmetatable(L, "fs.stdlibcpp.filesystem.recursive_directory_iterator");
	lua_setmetatable(L, -2);
	lua_pushcclosure(L, lua_fs_forall_dir_next, 1);
	return 1;
}

int lua_fs_foreach_dir_next(lua_State* L)
{
	int idx = lua_upvalueindex(1);
	luaL_argcheck(L, lua_type(L, idx) == LUA_TUSERDATA, 1, "Expected first argument to fs.foreach_dir_next to be userdata<std::filesystem::directory_iterator>");
	void* ud = luaL_checkudata(L, idx, "fs.stdlibcpp.filesystem.directory_iterator");
	luaL_argcheck(L, ud != nullptr, 1, "Expected first argument's userdata type to be a std::filesystem::directory_iterator");
	std::filesystem::directory_iterator& x = *(std::filesystem::directory_iterator*)ud;
	if (x._At_end())
		return 0;
	std::string p = (*x).path().string();
	lua_pushlstring(L, p.c_str(), p.size());
	x++;
	return 1;
}

int lua_fs_forall_dir_next(lua_State* L)
{
	int idx = lua_upvalueindex(1);
	luaL_argcheck(L, lua_type(L, idx) == LUA_TUSERDATA, 1, "Expected first argument to fs.forall_dir_next to be userdata<std::filesystem::recursive_directory_iterator>");
	void* ud = luaL_checkudata(L, idx, "fs.stdlibcpp.filesystem.recursive_directory_iterator");
	luaL_argcheck(L, ud != nullptr, 1, "Expected first argument's userdata type to be a std::filesystem::recursive_directory_iterator");
	std::filesystem::recursive_directory_iterator& x = *(std::filesystem::recursive_directory_iterator*)ud;
	if (x == std::filesystem::recursive_directory_iterator{})
		return 0;
	std::string p = (*x).path().string();
	lua_pushlstring(L, p.c_str(), p.size());
	x++;
	return 1;
}

int lua_fs_foreach_dir_dtor(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TUSERDATA, 1, "Expected first argument to fs.foreach_dir_next to be userdata<std::filesystem::directory_iterator>");
	void* ud = luaL_checkudata(L, 1, "fs.stdlibcpp.filesystem.directory_iterator");
	luaL_argcheck(L, ud != nullptr, 1, "Expected first argument's userdata type to be a std::filesystem::directory_iterator");
	std::filesystem::directory_iterator& x = *(std::filesystem::directory_iterator*)ud;
	x.~directory_iterator();
	return 0;
}

int lua_fs_forall_dir_dtor(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TUSERDATA, 1, "Expected first argument to fs.forall_dir_next to be userdata<std::filesystem::recursive_directory_iterator>");
	void* ud = luaL_checkudata(L, 1, "fs.stdlibcpp.filesystem.recursive_directory_iterator");
	luaL_argcheck(L, ud != nullptr, 1, "Expected first argument's userdata type to be a std::filesystem::recursive_directory_iterator");
	std::filesystem::recursive_directory_iterator& x = *(std::filesystem::recursive_directory_iterator*)ud;
	x.~recursive_directory_iterator();
	return 0;
}

int lua_fs_relative(lua_State* L)
{
	try {
		luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a path to make the second argument relative to");
		luaL_argcheck(L, lua_type(L, 2) == LUA_TSTRING, 2, "Expected second argument to be a string representing a full file path");
		const std::string f = std::filesystem::relative(std::filesystem::path(lua_tostring(L, 1)), std::filesystem::path(lua_tostring(L, 2))).string();
		lua_pushlstring(L, f.c_str(), f.size());
		return 1;
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
}

int lua_fs_absolute(lua_State* L)
{
	try {
		luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
		const std::string f = std::filesystem::absolute(std::filesystem::path(lua_tostring(L, 1))).string();
		lua_pushlstring(L, f.c_str(), f.size());
		return 1;
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
}

int lua_fs_file_name(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		const std::string f = std::filesystem::path(lua_tostring(L, 1)).filename().string();
		lua_pushlstring(L, f.c_str(), f.size());
		return 1;
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
}

int lua_fs_file_extension(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		const std::string f = std::filesystem::path(lua_tostring(L, 1)).extension().string();
		lua_pushlstring(L, f.c_str(), f.size());
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 1;
}

int lua_fs_file_directory(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		const std::string f = std::filesystem::path(lua_tostring(L, 1)).remove_filename().string();
		lua_pushlstring(L, f.c_str(), f.size());
		return 1;
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}

int lua_fs_is_newer(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	const char* test_file = lua_tostring(L, 1);
	if (!std::filesystem::exists(test_file))
	{
		lua_pushboolean(L, false);
		return 1;
	}

	if (lua_type(L, 2) == LUA_TTABLE)
	{
		size_t tbl_len = lua_objlen(L, 2);
		lua_newtable(L);
		for (size_t i = 1; i <= tbl_len; ++i)
		{
			lua_rawgeti(L, 2, i);
			if (lua_type(L, -1) != LUA_TSTRING) 
				return luaL_error(L, "Expected second argument at index %llu to be a string, not %s\n",
					i, lua_typename(L, lua_type(L, -1)));
			const char* target = lua_tostring(L, 2);
			lua_pushboolean(L,
				std::filesystem::last_write_time(test_file) > std::filesystem::last_write_time(target));
			lua_rawseti(L, -2, i);
			lua_pop(L, 1);
		}
		return 1;
	}
	else if (lua_type(L, 2) == LUA_TSTRING)
	{
		const char* target = lua_tostring(L, 2);
		if (!std::filesystem::exists(target))
		{
			lua_pushboolean(L, true);
			return 1;
		}
		lua_pushboolean(L,
			std::filesystem::last_write_time(test_file) > std::filesystem::last_write_time(target));
		return 1;
	}
	else return luaL_error(L, "Expected second argument to be a string or table of strings");
}

int lua_fs_filter_newer(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	const char* test_file = lua_tostring(L, 1);
	if (!std::filesystem::exists(test_file))
	{
		lua_pushnil(L);
		return 1;
	}

	if (lua_type(L, 2) == LUA_TTABLE)
	{
		size_t tbl_len = lua_objlen(L, 2);
		lua_newtable(L);
		for (size_t i = 1; i <= tbl_len; ++i)
		{
			lua_rawgeti(L, 2, i);
			if (lua_type(L, -1) != LUA_TSTRING)
				return luaL_error(L, "Expected second argument at index %llu to be a string, not %s\n",
					i, lua_typename(L, lua_type(L, -1)));
			const char* target = lua_tostring(L, 2);
			if (std::filesystem::last_write_time(test_file) > std::filesystem::last_write_time(target))
			{
				lua_pushvalue(L, -1);
				lua_rawseti(L, -2, lua_objlen(L, -2)+1);
			}
			lua_pop(L, 1);
		}
		return 1;
	}
	else if (lua_type(L, 2) == LUA_TSTRING)
	{
		const char* target = lua_tostring(L, 2);
		if (!std::filesystem::exists(target))
		{
			lua_pushvalue(L, 2);
			return 1;
		}
		
		if (std::filesystem::last_write_time(test_file) > std::filesystem::last_write_time(target))
			lua_pushvalue(L, 2);
		else lua_pushnil(L);

		return 1;
	}
	else return luaL_error(L, "Expected second argument to be a string or table of strings");
}

int lua_fs_last_modified(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		const uint64_t file_last_write = std::chrono::duration_cast<std::chrono::milliseconds>(
			std::chrono::clock_cast<std::chrono::system_clock>(
				std::filesystem::last_write_time(lua_tostring(L, 1))
			).time_since_epoch()
		).count();
		lua_pushinteger(L, file_last_write);
		return 1;
	} catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}

int lua_fs_exists(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		lua_pushboolean(L, std::filesystem::exists(lua_tostring(L, 1)));
		return 1;
	}
	catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}

int lua_fs_is_dir(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		lua_pushboolean(L, std::filesystem::is_directory(lua_tostring(L, 1)));
		return 1;
	}
	catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}

int lua_fs_is_file(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		lua_pushboolean(L, std::filesystem::is_regular_file(lua_tostring(L, 1)));
		return 1;
	}
	catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}

int lua_fs_create_dir(lua_State* L)
{
	luaL_argcheck(L, lua_type(L, 1) == LUA_TSTRING, 1, "Expected first argument to be a string representing a full file path");
	try {
		lua_pushboolean(L, std::filesystem::create_directories(lua_tostring(L, 1)));
		return 1;
	}
	catch (const std::exception& e)
	{
		lua_pushnil(L);
		lua_pushstring(L, e.what());
		return 2;
	}
	return 0;
}