#include <algorithm>
#include <fcntl.h>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <sys/stat.h>
#include <unistd.h>
#include <unordered_map>
#include <vector>

#define PROC_LOCATION "/proc"

namespace Animia::Linux {

std::vector<int> get_all_pids() {
	std::vector<int> ret;
	for (const auto& dir : std::filesystem::directory_iterator(PROC_LOCATION)) {
		int pid;
		try {
			pid = std::stoi(dir.path().stem());
		} catch (std::invalid_argument) {
			continue;
		}
		ret.push_back(pid);
	}
	return ret;
}

std::string get_process_name(int pid) {
	std::string path = PROC_LOCATION "/" + std::to_string(pid) + "/comm";
	std::ifstream t(path);
	std::stringstream buf;
	buf << t.rdbuf();

	std::string str = buf.str();
	str.erase(std::remove(str.begin(), str.end(), '\n'), str.end());
	return str;
}

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

static bool are_flags_ok(int pid, int fd) {
	std::string path = PROC_LOCATION "/" + std::to_string(pid) + "/fdinfo/" + std::to_string(fd);
	std::ifstream t(path);
	std::stringstream buffer;
	buffer << t.rdbuf();
	std::string raw;
	int flags = 0;
	while (std::getline(buffer, raw)) {
		if (raw.rfind("flags:", 0) == 0) {
			flags = std::stoi(raw.substr(raw.find_last_not_of("0123456789") + 1));
		}
	}
	if (flags & O_WRONLY || flags & O_RDWR)
		return false;
	return true;
}

static int get_size_of_link(std::string link) {
	struct stat sb;
	if (lstat(link.c_str(), &sb) == -1)
		return -1;
	return sb.st_size + 1;
}

std::vector<std::string> get_open_files(int pid) {
	std::vector<std::string> ret;
	std::string path = PROC_LOCATION "/" + std::to_string(pid) + "/fd";
	try {
		for (const auto& dir : std::filesystem::directory_iterator(path)) {
			if (!are_flags_ok(pid, std::stoi(dir.path().stem())))
				continue;
			/* get filename size */
			int size = get_size_of_link(dir.path().string());
			/* allocate buffer */
			std::string buf(size, '\0');
			// std::string buf('\0', size);
			/* read filename to buffer */
			int r = readlink(dir.path().c_str(), &buf.front(), buf.length());

			if (r < 0)
				continue;
			if (!is_regular_file(buf))
				continue;
			ret.push_back(buf);
		}
	} catch (std::filesystem::filesystem_error const& ex) {
		if (ex.code().value() != 13) // 13 == permissions error, common with /proc, ignore
			throw;
	}
	return ret;
}

std::unordered_map<int, std::vector<std::string>> get_all_open_files() {
	std::unordered_map<int, std::vector<std::string>> map;
	std::vector<int> pids = get_all_pids();
	for (int i : pids)
		map[i] = get_open_files(i);
	return map;
}

} // namespace Animia::Linux