#include "bounded_thread_pool.hpp"

namespace fz
{

struct bounded_async_task::impl: intrusive_node
{
	async_task handle{};

	bounded_thread_pool *pool{};
	std::function<void()> task{};
	std::function<void()> on_timeout{};
	monotonic_clock expiration{};
};

bounded_thread_pool::bounded_thread_pool(thread_pool &pool, std::size_t max_number_of_threads, duration timeout)
	: pool_(pool)
	, number_of_available_threads_(max_number_of_threads)
{
	if (timeout.get_milliseconds() > 0) {
		timeout_ = std::move(timeout);
	}
}

bounded_thread_pool::~bounded_thread_pool()
{
	scoped_lock l(m_);
	quit_ = true;
	spooler_cond_.signal(l);
}

bounded_async_task bounded_thread_pool::spawn(std::function<void ()> task, std::function<void ()> on_timeout)
{
	if (!task) {
		return {};
	}

	async_task handle;
	scoped_lock l(m_);

	if (auto state = try_to_spawn(std::move(task), handle); state == spawn_not_attempted) {
		if (!spooler_) {
			spooler_ = pool_.spawn([this]{task_spooler();});
			if (!spooler_) {
				return {};
			}
		}

		bounded_async_task ret;
		ret.impl_ = std::unique_ptr<bounded_async_task::impl>(new bounded_async_task::impl{{}, {}, this, std::move(task), std::move(on_timeout), monotonic_clock::now() + timeout_});

		if (waiting_.empty()) {
			spooler_cond_.signal(l);
		}

		waiting_.push_back(*ret.impl_);
		return ret;
	}
	else
	if (state != spawn_failed) {
		bounded_async_task ret;
		ret.impl_ = std::unique_ptr<bounded_async_task::impl>(new bounded_async_task::impl{{}, std::move(handle)});
		return ret;
	}

	return {};
}

bounded_thread_pool::spawn_state bounded_thread_pool::try_to_spawn(std::function<void ()> &&task, async_task &handle)
{
	if (number_of_available_threads_ > 0) {
		handle.join();

		handle = pool_.spawn([this, task = std::move(task)] {
			task();

			scoped_lock l(m_);
			if (!quit_) {
				if (number_of_available_threads_ == 0 && !waiting_.empty()) {
					spooler_cond_.signal(l);
				}

				number_of_available_threads_ += 1;
			}
		});

		if (handle) {
			number_of_available_threads_ -= 1;
			return spawn_succeeded;
		}

		return spawn_failed;
	}

	return spawn_not_attempted;
}

void bounded_thread_pool::task_spooler()
{
	scoped_lock l(m_);

	if (timeout_) {
		while (!quit_) {
			if (waiting_.empty()) {
				spooler_cond_.wait(l);
			}
			else {
				auto &wt = waiting_.front();

				auto now = monotonic_clock::now();

				if (now >= wt.expiration) {
					auto on_timeout = std::move(wt.on_timeout);
					waiting_.pop_front();

					if (on_timeout) {
						l.unlock();
						on_timeout();
						l.lock();
					}
				}
				else
				if (auto state = try_to_spawn(std::move(wt.task), wt.handle); state == spawn_succeeded) {
					waiting_.pop_front();
				}
				else
				if (state == spawn_failed) {
					// Reuse the timeout handler as an error handler.
					auto on_timeout = std::move(wt.on_timeout);
					waiting_.pop_front();

					if (on_timeout) {
						l.unlock();
						on_timeout();
						l.lock();
					}
				}
				else {
					spooler_cond_.wait(l, wt.expiration - now);
				}
			}
		}
	}
	else {
		while (!quit_) {
			if (waiting_.empty()) {
				spooler_cond_.wait(l);
			}
			else {
				auto &wt = waiting_.front();

				if (auto state = try_to_spawn(std::move(wt.task), wt.handle); state == spawn_succeeded) {
					waiting_.pop_front();
				}
				else
				if (state == spawn_failed) {
					// Reuse the timeout handler as an error handler.
					auto on_timeout = std::move(wt.on_timeout);
					waiting_.pop_front();

					if (on_timeout) {
						l.unlock();
						on_timeout();
						l.lock();
					}
				}
				else {
					spooler_cond_.wait(l);
				}
			}
		}
	}
}

bounded_async_task::bounded_async_task() noexcept
{
};

bounded_async_task::~bounded_async_task() noexcept
{
	cancel();
}

void swap(bounded_async_task &lhs, bounded_async_task &rhs) noexcept
{
	using std::swap;

	swap(lhs.impl_, rhs.impl_);
}

bounded_async_task::bounded_async_task(bounded_async_task &&other) noexcept
	: bounded_async_task()
{
	swap(*this, other);
}

bounded_async_task &bounded_async_task::operator=(bounded_async_task &&other) noexcept
{
	bounded_async_task copy(std::move(other));

	swap(*this, copy);

	return *this;
}

bool bounded_async_task::cancel() noexcept
{
	if (impl_) {
		if (auto *p = impl_->pool) {
			scoped_lock l(p->m_);

			if (!impl_->is_removed()) {
				impl_->remove();
				return true;
			}
		}
	}

	return false;
}

}
