#include "animia/fd/proc.h"
#include "animia.h"
#include "animia/util.h"

#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>

#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

static constexpr std::string_view PROC_LOCATION = "/proc";

namespace animia::internal::proc {

template<typename T = int>
static T StringToInt(const std::string& str, T def = 0) {
	std::istringstream s(str);
	s >> def;
	return def;
}

static bool IsRegularFile(std::string link) {
	struct stat sb;
	if (stat(link.c_str(), &sb) == -1)
		return false;

	return S_ISREG(sb.st_mode);
}

static bool AreFlagsOk(pid_t pid, int fd) {
	const std::filesystem::path path = std::filesystem::path(PROC_LOCATION) / std::to_string(pid) / "fdinfo" / std::to_string(fd);

	std::ifstream file(path);
	if (!file)
		return false;

	int flags = 0;
	for (std::string line; std::getline(file, line); )
		if (line.find("flags:", 0) == 0)
			flags = StringToInt(line.substr(line.find_last_not_of("0123456789") + 1));

	if (flags & O_WRONLY || flags & O_RDWR)
		return false;

	return true;
}

static bool GetFilenameFromFd(std::string link, std::string& out) {
	/* gets around stupid linux limitation where /proc doesn't
	 * give actual size readings of the string
	*/
	constexpr size_t OUT_MAX = (1ul << 15); // 32KiB
	out.resize(32);

	ssize_t exe_used = 0;
	do {
		out.resize(out.length() * 2);

		exe_used = readlink(link.c_str(), &out.front(), out.length());
		if (exe_used == (ssize_t)-1 || exe_used < (ssize_t)1)
			return false; // we got a bad result. SAD!
	} while (out.length() < OUT_MAX && exe_used >= static_cast<ssize_t>(out.length()));

	out.resize(out.find('\0'));

	return true;
}

static bool GetProcessName(pid_t pid, std::string& result) {
	const std::filesystem::path path = std::filesystem::path(PROC_LOCATION) / std::to_string(pid) / "comm";

	if (!util::ReadFile(path, result))
		return false;

	result.erase(std::remove(result.begin(), result.end(), '\n'), result.end());
	return true;
}

bool EnumerateOpenProcesses(process_proc_t process_proc) {
	bool success = false;

	for (const auto& dir : std::filesystem::directory_iterator{PROC_LOCATION}) {
		Process proc;

		try {
			proc.pid = StringToInt(dir.path().stem());
			success = true;
		} catch (std::invalid_argument const& ex) {
			continue;
		}

		if (!GetProcessName(proc.pid, proc.name))
			continue;

		if (!process_proc(proc))
			return false;
	}

	return success;
}

bool EnumerateOpenFiles(const std::set<pid_t>& pids, open_file_proc_t open_file_proc) {
	if (!open_file_proc)
		return false;

	for (const auto& pid : pids) {
		const std::filesystem::path path = std::filesystem::path(PROC_LOCATION) / std::to_string(pid) / "fd";

		for (const auto& dir : std::filesystem::directory_iterator{path}) {
			if (!AreFlagsOk(pid, StringToInt(dir.path().stem())))
				continue;

			std::string name;
			if (!GetFilenameFromFd(dir.path(), name))
				continue;

			if (!IsRegularFile(name))
				continue;

			if (!open_file_proc({pid, name}))
				return false;
		}
	}
	return true;
}

} // namespace animia::internal::linux