#include "animone/fd/bsd.h"
#include "animone.h"
#include "animone/fd.h"

#include <sys/file.h>
#include <sys/filedesc.h>
#include <sys/param.h>
#include <sys/queue.h>
#include <sys/sysctl.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/vnode.h>

#ifdef HAVE_KVM_GETFILES
#	include <kvm.h>
#elif defined(LIBUTIL)
#	include <libutil.h>
#endif

#include <string>
#include <iostream>

namespace animone::internal::bsd {

static std::string Basename(const std::string& name) {
	size_t s = name.find_last_of('/');

	if (s == std::string::npos)
		return name;

	return name.substr(s, name.size());
}

bool GetProcessName(pid_t pid, std::string& name) {
#ifdef HAVE_KVM_GETFILES
	char errbuf[_POSIX2_LINE_MAX];
	kvm_t* kernel = kvm_openfiles(NULL, NULL, NULL, O_RDONLY, errbuf);
	if (!kernel)
		return false;

	int entries = 0;
	struct kinfo_proc* kinfo = kvm_getprocs(kernel, KERN_PROC_PID, pid, &entries);
	if (!kinfo) {
		kvm_close(kernel);
		return false;
	}

	if (entries < 1) {
		kvm_close(kernel);
		return false;
	}

	name = Basename(kinfo[0].ki_paddr->p_comm);

	return true;
#else
	/* use sysctl as a fallback */
	static const int mib[] = {CTL_KERN, KERN_PROC, KERN_PROC_PID, pid};

	struct kinfo_proc result;

	size_t length = 1;
	if (sysctl((int*)mib, (sizeof(mib) / sizeof(*mib)) - 1, &result, &length, NULL, 0) == -1)
		return false;

	name = Basename(result.ki_comm);

	return true;
#endif
}

/* Most of the BSDs share the common kvm library,
 * so accessing this information can be trivial.
 */
bool EnumerateOpenProcesses(process_proc_t process_proc) {
#ifdef HAVE_KVM_GETFILES
	char errbuf[_POSIX2_LINE_MAX];
	kvm_t* kernel = kvm_openfiles(NULL, NULL, NULL, O_RDONLY, errbuf);
	if (!kernel)
		return false;

	int entries = 0;
	struct kinfo_proc* kinfo = kvm_getprocs(kernel, KERN_PROC_ALL, 0, &entries);
	if (!kinfo) {
		kvm_close(kernel);
		return false;
	}

	for (int i = 0; i < entries; i++) {
		if (!process_proc({kinfo[i].ki_paddr->p_pid, Basename(kinfo[i].ki_paddr->p_comm)})) {
			kvm_close(kernel);
			return false;
		}
	}

	kvm_close(kernel);

	return true;
#else
	/* use sysctl as a fallback */
	static const int mib[] = {CTL_KERN, KERN_PROC, KERN_PROC_ALL, 0};
	size_t length = 0;

	sysctl((int*)mib, (sizeof(mib) / sizeof(*mib)) - 1, NULL, &length, NULL, 0);

	std::unique_ptr<struct kinfo_proc[]> result;
	result.reset(new struct kinfo_proc[length]);

	if (!result.get())
		return false;

	/* actually get our results */
	if (sysctl((const int*)mib, (sizeof(mib) / sizeof(*mib)) - 1, result.get(), &length, NULL, 0) == ENOMEM) {
		result.reset();
		throw std::bad_alloc();
	}

	if (length < sizeof(struct kinfo_proc))
		return false;

	for (int i = 0; i < length / sizeof(result[0]); i++)
		if (!process_proc({result[i].ki_pid, result[i].ki_comm}))
			return false;

	return true;
#endif
}

bool EnumerateOpenFiles(const std::set<pid_t>& pids, open_file_proc_t open_file_proc) {
#ifdef HAVE_KVM_GETFILES
	char errbuf[_POSIX2_LINE_MAX];
	kvm_t* kernel = kvm_openfiles(NULL, NULL, NULL, O_RDONLY, errbuf);
	if (!kernel)
		return false;

	for (const auto& pid : pids) {
		int cnt;
		struct kinfo_file* kfile = kvm_getfiles(kernel, KERN_FILE_BYPID, pid, &cnt);
		if (!kfile) {
			kvm_close(kernel);
			return false;
		}

		for (int i = 0; i < cnt; i++) {
			if (!open_file_proc({pid, kfile[i].kf_path})) {
				kvm_close(kernel);
				return false;
			}
		}
	}

	kvm_close(kernel);

	return true;
#elif defined(LIBUTIL)
	/* does this code even work? */
	for (const auto& pid : pids) {
		int cnt;
		std::unique_ptr<struct kinfo_file[]> files(kinfo_getfile(pid, &cnt));
		if (!files)
			return false;

		for (int i = 0; i < cnt; i++) {
			const struct kinfo_file& current = files[i];
			if (current.kf_vnode_type != KF_VTYPE_VREG)
				continue;

			if (!open_file_proc({pid, current.kf_path}))
				return false;
		}
	}

	return true;
#else
	for (const auto& pid : pids) {
		int mib[6] = {CTL_KERN, KERN_FILE2, KERN_FILE_BYPID, pid, sizeof(struct kinfo_file), 0};

		size_t len = 0;
		if (sysctl(mib, sizeof(mib) / sizeof(mib[0]), NULL, &len, NULL, 0) == -1)
			return false;

		mib[5] = len / sizeof(struct kinfo_file);

		std::unique_ptr<struct kinfo_file[]> buf(new struct kinfo_file[mib[5]]);
		if (!buf)
			return false;

		if (sysctl(mib, sizeof(mib) / sizeof(mib[0]), buf.get(), &len, NULL, 0) == -1)
			return false;

		/* TODO: check kfile[i].ki_ofileflags */
		for (size_t i = 0; i < mib[5]; i++)
			if (!open_file_proc({pid, kfile[i].kf_path}))
				return false;
	}

	return true;
#endif
}

} // namespace animone::internal::kvm