#include "ThreadPinning.hpp"

#include <algorithm>
#include <fstream>
#include <set>
#include <sstream>
#include <string>

#if defined(__linux__) || defined(__ANDROID__)
#include <sched.h>
#define TP_THREAD_PINNING_SUPPORTED
#endif

namespace tp {

	namespace {
		std::vector<int> parseCpuList(const std::string& text) {
			std::vector<int> cpus;
			std::stringstream stream(text);
			std::string range;
			while (std::getline(stream, range, ',')) {
				if (range.empty()) continue;
				const size_t dash = range.find('-');
				if (dash == std::string::npos) {
					cpus.push_back(std::stoi(range));
					continue;
				}
				const int first = std::stoi(range.substr(0, dash));
				const std::string tail = range.substr(dash + 1);
				const int last = tail.empty() ? first : std::stoi(tail);
				for (int cpu = first; cpu <= last; cpu++)
					cpus.push_back(cpu);
			}
			return cpus;
		}

		std::vector<int> siblingsOf(int cpu) {
			std::ifstream file("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/thread_siblings_list");
			std::string text;
			if (!file || !std::getline(file, text)) return { cpu };
			const std::vector<int> siblings = parseCpuList(text);
			return siblings.empty() ? std::vector<int>{ cpu } : siblings;
		}
	}

	std::vector<int> ThreadPinning::allowedCpus() {
#ifdef TP_THREAD_PINNING_SUPPORTED
		cpu_set_t set;
		CPU_ZERO(&set);
		if (sched_getaffinity(0, sizeof(set), &set) != 0) return {};
		std::vector<int> cpus;
		for (int cpu = 0; cpu < CPU_SETSIZE; cpu++)
			if (CPU_ISSET(cpu, &set)) cpus.push_back(cpu);
		return cpus;
#else
		return {};
#endif
	}

	std::vector<std::vector<int>> ThreadPinning::siblingGroups(const std::vector<int>& allowed) {
		const std::vector<int> cpus = allowed.empty() ? allowedCpus() : allowed;
		std::vector<std::vector<int>> groups;
		std::set<int> seen;
		for (int cpu : cpus) {
			if (seen.count(cpu)) continue;
			std::vector<int> group;
			for (int sibling : siblingsOf(cpu)) {
				if (!std::count(cpus.begin(), cpus.end(), sibling)) continue;
				group.push_back(sibling);
				seen.insert(sibling);
			}
			if (group.empty()) group.push_back(cpu);
			groups.push_back(group);
		}
		return groups;
	}

	void ThreadPinning::pinCallingThread(const std::vector<int>& cpus) {
#ifdef TP_THREAD_PINNING_SUPPORTED
		if (cpus.empty()) return;
		cpu_set_t set;
		CPU_ZERO(&set);
		for (int cpu : cpus)
			CPU_SET(cpu, &set);
		sched_setaffinity(0, sizeof(set), &set);
#endif
	}
}