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

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

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

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

namespace animone::internal::proc {

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 = util::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) {
	/* /proc is a "virtual filesystem", so we have to guess the path size. yippee! */
	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 IsSystemFile(const std::string& path) {
	static constexpr std::array<std::string_view, 9> invalid_paths = {"/boot", "/dev", "/bin", "/usr", "/opt",
	                                                                  "/proc", "/var", "/etc", "/dev"};

	for (const auto& invalid_path : invalid_paths) {
		if (!path.rfind(invalid_path, 0)) {
			return true;
		}
	}

	return false;
}

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 = util::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, util::StringToInt(dir.path().stem())))
				continue;

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

			if (!IsRegularFile(name))
				continue;

			if (IsSystemFile(name))
				continue;

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

} // namespace animia::internal::proc
