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

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

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

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

namespace animia::internal::proc {

/* this uses dirent instead of std::filesystem; it would make a bit
   more sense to use the latter, but this is platform dependent already :) */
static std::vector<std::string> GetAllFilesInDir(const std::string& _dir) {
	std::vector<std::string> ret;

	DIR* dir = opendir(_dir.c_str());
	if (!dir)
		return ret;

	struct dirent* dp;
	while ((dp = readdir(dir)) != NULL) {
		if (!(!strcmp(dp->d_name, ".") || !strcmp(dp->d_name, "..")))
			ret.push_back(_dir + "/" + dp->d_name);
	}

	closedir(dir);
	return ret;
}

static std::string Basename(const std::string& path) {
	return path.substr(path.find_last_of("/") + 1, path.length());
}

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::string path = std::string(PROC_LOCATION) + "/" + std::to_string(pid) + "/fdinfo/" + std::to_string(fd);

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

	int flags = 0;
	for (std::string line; std::getline(file, line); )
		if (line.find("flags:", 0) == 0)
			flags = std::stoi(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 = (1 << 15); // 32KiB
	out.resize(1024);

	for (ssize_t exe_used = 0;
	     out.length() < OUT_MAX && exe_used >= (ssize_t)(out.length() - 1);
		 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, i think
	}

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

	return true;
}

static std::string GetProcessName(pid_t pid) {
	std::string result;

	const std::string path = std::string(PROC_LOCATION) + "/" + std::to_string(pid) + "/comm";

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

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

	return result;
}

bool ProcFdTools::EnumerateOpenProcesses(process_proc_t process_proc) {
	bool success = false;
	for (const auto& dir : GetAllFilesInDir(std::string(PROC_LOCATION))) {
		pid_t pid;
		try {
			pid = std::stoul(Basename(dir));
			success = true;
		} catch (std::invalid_argument) {
			continue;
		}
		if (!process_proc({pid, GetProcessName(pid)}))
			return false;
	}
	return success;
}

bool ProcFdTools::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::string path = std::string(PROC_LOCATION) + "/" + std::to_string(pid) + "/fd";

		for (const auto& dir : GetAllFilesInDir(path)) {
			if (!AreFlagsOk(pid, std::stoi(Basename(dir))))
				continue;

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

			if (!IsRegularFile(name))
				continue;

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

} // namespace animia::internal::linux