//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

module Crafter.Thread:ThreadPool_impl;
import std;
import :ThreadPool;

using namespace Crafter;


void ThreadPool::Start(int maxThreads)
{
	stopped = false;
	if (maxThreads > 0) {
		threadCount = maxThreads;
	}
	else {
		threadCount = std::thread::hardware_concurrency();
	}
	threadStorage = new ThreadStorage[threadCount];
	for (unsigned int i = 0; i < threadCount - 1; i++)
	{
		threadStorage[i].thread = std::thread(ThreadPool::WaitAndExecuteTask, threadStorage + i);
		threadIdMap[threadStorage[i].thread.get_id()] = &threadStorage[i];
	}
	threadIdMap[std::this_thread::get_id()] = &threadStorage[threadCount - 1];
	activeThreads = threadCount - 1;
}

void ThreadPool::Stop()
{
	// This pool's own: a JoinPool caller may Start another as soon as it
	// returns, replacing threadStorage.
	ThreadStorage* const storage = threadStorage;
	const unsigned int count = threadCount;
	if (storage == nullptr) return;
	stopped = true;
	// Taking the slot's mutex orders the store above before the next check
	// of `stopped` by whoever waits on the slot: without it one between
	// checking the condition and going to sleep misses the notify, and never
	// wakes.
	auto wake = [](ThreadStorage& slot) {
		{
			std::lock_guard lk(slot.mutex);
		}
		slot.cv.notify_one();
	};
	for (unsigned int i = 0; i < count - 1; i++) wake(storage[i]);
	for (unsigned int i = 0; i < count - 1; i++)
	{
		if (storage[i].thread.joinable()) storage[i].thread.join();
	}
	// The JoinPool caller's slot last: a JoinPool waiting there returns once
	// the workers have stopped (one busy with a task sees `stopped` itself
	// when it's done).
	wake(storage[count - 1]);
}

void ThreadPool::Enqueue(std::function<void(void)> task)
{
	unsigned int min = std::numeric_limits<unsigned int>::max();
	unsigned int minIndex = 0;
	for (unsigned int i = 0; i < activeThreads; i++)
	{
		unsigned int taskCount = threadStorage[i].taskCount;
		if (taskCount == 0) {
			minIndex = i;
			break;
		}
		if (taskCount < min) {
			min = taskCount;
			minIndex = i;
		}
	}
	{
		std::lock_guard lk(threadStorage[minIndex].mutex);
		threadStorage[minIndex].buffer.push_back(task);
		threadStorage[minIndex].taskCount++;
	}
	threadStorage[minIndex].cv.notify_one();
}

void ThreadPool::Enqueue(std::vector<std::function<void(void)>> tasks)
{
	const unsigned int workers = activeThreads;
	if (workers == 0) {
		// Start(1) runs no worker threads; like a single Enqueue, the batch
		// waits for the caller's JoinPool.
		{
			std::lock_guard lk(threadStorage[0].mutex);
			threadStorage[0].taskCount += static_cast<unsigned int>(tasks.size());
			threadStorage[0].buffer.insert(threadStorage[0].buffer.end(),
				std::make_move_iterator(tasks.begin()), std::make_move_iterator(tasks.end()));
		}
		threadStorage[0].cv.notify_one();
		return;
	}

	// As even a split as the count allows: the first `remainder` workers
	// take one task more. Every task goes to exactly one worker, and each
	// worker's taskCount grows by what it was given, which the single-task
	// Enqueue relies on to pick the least loaded worker.
	const std::size_t perWorker = tasks.size() / workers;
	const std::size_t remainder = tasks.size() % workers;
	std::size_t first = 0;
	for (unsigned int i = 0; i < workers && first < tasks.size(); i++)
	{
		const std::size_t count = perWorker + (i < remainder ? 1 : 0);
		{
			std::lock_guard lk(threadStorage[i].mutex);
			threadStorage[i].taskCount += static_cast<unsigned int>(count);
			threadStorage[i].buffer.insert(threadStorage[i].buffer.end(),
				std::make_move_iterator(tasks.begin() + static_cast<std::ptrdiff_t>(first)),
				std::make_move_iterator(tasks.begin() + static_cast<std::ptrdiff_t>(first + count)));
		}
		threadStorage[i].cv.notify_one();
		first += count;
	}
}

void ThreadPool::JoinPool()
{
	activeThreads++;
	ThreadPool::WaitAndExecuteTask(threadStorage + threadCount - 1);
}

void ThreadPool::WaitAndExecuteTask(ThreadStorage* threadStorage)
{
	while (true) {
		std::unique_lock<std::mutex> lock(threadStorage->mutex);
		threadStorage->cv.wait(lock, [threadStorage] { return !threadStorage->buffer.empty() || stopped; });
		// The buffer is only read under the lock: the loop used to test it
		// unlocked while Enqueue appended to it. Stopping still waits for
		// everything queued to run.
		if (threadStorage->buffer.empty()) {
			return;
		}
		threadStorage->buffer.swap(threadStorage->tasks);
		lock.unlock();
		for (std::function<void(void)>& task : threadStorage->tasks)
		{
			task();
			threadStorage->taskCount--;
		}
		threadStorage->tasks.clear();
	}
}