#pragma once

#include <psemek/util/function.hpp>
#include <psemek/util/exception.hpp>

#include <atomic>
#include <mutex>
#include <condition_variable>
#include <exception>
#include <optional>
#include <vector>

namespace psemek::async
{

	namespace detail
	{

		template <typename T>
		struct task_value_container
		{
			using type = std::optional<T>;
		};

		template <>
		struct task_value_container<void>
		{
			using type = bool;
		};

		template <typename T>
		struct then_function
		{
			using type = util::function<void(T &)>;
		};

		template <>
		struct then_function<void>
		{
			using type = util::function<void()>;
		};

		template <typename T>
		struct future_result_type
		{
			using type = T &;
		};

		template <>
		struct future_result_type<void>
		{
			using type = void;
		};

		struct cancel_token
		{};

		template <typename T>
		struct task_state
		{
			std::mutex value_mutex;
			typename task_value_container<T>::type value{};
			std::exception_ptr exception;
			std::condition_variable value_cv;

			std::shared_ptr<cancel_token> shared_cancel;
			std::weak_ptr<cancel_token> weak_cancel;

			std::mutex then_mutex;
			std::vector<typename then_function<T>::type> then_funcs;
			std::vector<util::function<void(std::exception_ptr)>> then_exception_funcs;
		};

	}

	struct empty_future_error
		: util::exception
	{
		empty_future_error(util::stacktrace stacktrace = {})
			: util::exception("Future is empty", std::move(stacktrace))
		{}
	};

	struct canceled_task_error
		: util::exception
	{
		canceled_task_error(util::stacktrace stacktrace = {})
			: util::exception("Task canceled", std::move(stacktrace))
		{}
	};

	struct auto_cancel_tag{};
	constexpr auto_cancel_tag auto_cancel;

	template <typename T>
	struct future
	{
		using result_type = T;

		future() = default;
		future(future const &) = default;
		future(future &&) = default;

		future(std::shared_ptr<detail::task_state<T>> state, std::shared_ptr<detail::cancel_token> cancel)
			: state_(std::move(state))
			, cancel_(std::move(cancel))
		{}

		future & operator = (future const &) = default;
		future & operator = (future &&) = default;

		~future()
		{
			reset();
		}

		void reset()
		{
			state_.reset();
			cancel_.reset();
		}

		explicit operator bool() const
		{
			return static_cast<bool>(state_);
		}

		bool wait() const
		{
			if (!state_) throw empty_future_error{};
			std::unique_lock lock(state_->value_mutex);
			state_->value_cv.wait(lock, [this]{ return has_value_unsafe(); });
			return true;
		}

		template <typename Duration>
		bool wait_for(Duration period) const
		{
			if (!state_) throw empty_future_error{};
			std::unique_lock lock(state_->value_mutex);
			state_->value_cv.wait_for(lock, period, [this]{ return has_value_unsafe(); });
			return has_value_unsafe();
		}

		template <typename TimePoint>
		bool wait_until(TimePoint time) const
		{
			if (!state_) throw empty_future_error{};
			std::unique_lock lock(state_->value_mutex);
			state_->value_cv.wait_until(lock, time, [this]{ return has_value_unsafe(); });
			return has_value_unsafe();
		}

		typename detail::future_result_type<T>::type get()
		{
			if (!state_)
				throw empty_future_error{};
			if (!state_->weak_cancel.lock())
				throw canceled_task_error{};
			wait();
			if (state_->value)
			{
				if constexpr (std::is_same_v<T, void>)
					return;
				else
					return *(state_->value);
			}
			else
				std::rethrow_exception(state_->exception);
		}

		void cancel()
		{
			if (!state_) throw empty_future_error{};

			cancel_.reset();
		}

		bool ready() const
		{
			if (!state_) throw empty_future_error{};
			std::unique_lock lock(state_->value_mutex);
			return has_value_unsafe();
		}

		template <typename F>
		auto then(F && f);

	private:
		std::shared_ptr<detail::task_state<T>> state_;
		std::shared_ptr<detail::cancel_token> cancel_;

		bool has_value_unsafe() const
		{
			return state_->value || state_->exception;
		}
	};

	struct empty_promise_error
		: util::exception
	{
		empty_promise_error(util::stacktrace stacktrace = {})
			: util::exception("Promise is empty", std::move(stacktrace))
		{}
	};

	struct satisfied_promise_error
		: util::exception
	{
		satisfied_promise_error(util::stacktrace stacktrace = {})
			: util::exception("Promise already contains a value or exception", std::move(stacktrace))
		{}
	};

	template <typename T>
	struct promise
	{
		promise(std::nullptr_t)
		{}

		promise()
			: state_{std::make_shared<detail::task_state<T>>()}
			, cancel_{std::make_shared<detail::cancel_token>()}
		{
			state_->shared_cancel = cancel_;
			state_->weak_cancel = cancel_;
		}

		promise(auto_cancel_tag)
			: state_{std::make_shared<detail::task_state<T>>()}
			, cancel_{std::make_shared<detail::cancel_token>()}
		{
			state_->weak_cancel = cancel_;
		}

		promise(std::shared_ptr<detail::cancel_token> cancel)
			: state_{std::make_shared<detail::task_state<T>>()}
			, cancel_{cancel}
		{
			state_->weak_cancel = cancel_;
		}

		explicit operator bool() const
		{
			return static_cast<bool>(state_);
		}

		void set_value(T const & value)
		{
			if (!state_) throw empty_promise_error{};
			{
				std::lock_guard lock{state_->value_mutex};
				if (state_->value || state_->exception) throw satisfied_promise_error{};
				state_->value = value;
			}
			state_->value_cv.notify_all();

			std::lock_guard lock{state_->then_mutex};
			for (auto & func : state_->then_funcs)
				func(*state_->value);
		}

		void set_value(T && value)
		{
			if (!state_) throw empty_promise_error{};
			{
				std::lock_guard lock{state_->value_mutex};
				if (state_->value || state_->exception) throw satisfied_promise_error{};
				state_->value = std::move(value);
			}
			state_->value_cv.notify_all();

			std::lock_guard lock{state_->then_mutex};
			for (auto & func : state_->then_funcs)
				func(*state_->value);
		}

		void set_exception(std::exception_ptr e)
		{
			if (!state_) throw empty_promise_error{};
			{
				std::lock_guard lock{state_->value_mutex};
				if (state_->value || state_->exception) throw satisfied_promise_error{};
				state_->exception = std::move(e);

				std::lock_guard lock_then{state_->then_mutex};
				for (auto & func : state_->then_exception_funcs)
					func(state_->exception);
			}
			state_->value_cv.notify_all();
		}

		future<T> get_future()
		{
			return future<T>(state_, std::move(cancel_));
		}

		bool canceled() const
		{
			return !state_->weak_cancel.lock();
		}

	private:
		std::shared_ptr<detail::task_state<T>> state_;
		std::shared_ptr<detail::cancel_token> cancel_;
	};

	template <>
	struct promise<void>
	{
		promise(std::nullptr_t)
		{}

		promise()
			: state_{std::make_shared<detail::task_state<void>>()}
			, cancel_{std::make_shared<detail::cancel_token>()}
		{
			state_->shared_cancel = cancel_;
			state_->weak_cancel = cancel_;
		}

		promise(auto_cancel_tag)
			: state_{std::make_shared<detail::task_state<void>>()}
			, cancel_{std::make_shared<detail::cancel_token>()}
		{
			state_->weak_cancel = cancel_;
		}

		promise(std::shared_ptr<detail::cancel_token> cancel)
			: state_{std::make_shared<detail::task_state<void>>()}
			, cancel_{cancel}
		{
			state_->weak_cancel = cancel_;
		}

		explicit operator bool() const
		{
			return static_cast<bool>(state_);
		}

		void set_value()
		{
			if (!state_) throw empty_promise_error{};
			{
				std::lock_guard lock{state_->value_mutex};
				if (state_->value || state_->exception) throw satisfied_promise_error{};
				state_->value = true;
			}
			state_->value_cv.notify_all();

			std::lock_guard lock{state_->then_mutex};
			for (auto & func : state_->then_funcs)
				func();
		}

		void set_exception(std::exception_ptr e)
		{
			if (!state_) throw empty_promise_error{};
			{
				std::lock_guard lock{state_->value_mutex};
				if (state_->value || state_->exception) throw satisfied_promise_error{};
				state_->exception = std::move(e);

				std::lock_guard lock_then{state_->then_mutex};
				for (auto & func : state_->then_exception_funcs)
					func(state_->exception);
			}
			state_->value_cv.notify_all();
		}

		future<void> get_future()
		{
			return future<void>(state_, std::move(cancel_));
		}

		bool canceled() const
		{
			return !state_->weak_cancel.lock();
		}

	private:
		std::shared_ptr<detail::task_state<void>> state_;
		std::shared_ptr<detail::cancel_token> cancel_;
	};

	template <typename T>
	future<T> make_ready_future(T && x)
	{
		promise<T> p;
		p.set_value(std::move(x));
		return p.get_future();
	}

	template <typename T>
	future<T> make_ready_future(T const & x)
	{
		promise<T> p;
		p.set_value(x);
		return p.get_future();
	}

	template <typename Signature>
	struct packaged_task;

	template <typename R, typename ... Args>
	struct packaged_task<R(Args...)>
	{
		packaged_task()
			: promise_(nullptr)
		{}

		template <typename F>
		packaged_task(F && f)
			: func_(std::forward<F>(f))
		{}

		template <typename F>
		packaged_task(auto_cancel_tag tag, F && f)
			: promise_(tag)
			, func_(std::forward<F>(f))
		{}

		template <typename F>
		packaged_task(std::shared_ptr<detail::cancel_token> cancel, F && f)
			: promise_(cancel)
			, func_(std::forward<F>(f))
		{}

		explicit operator bool() const
		{
			return static_cast<bool>(promise_);
		}

		future<R> get_future()
		{
			return promise_.get_future();
		}

		template <typename ... Args1>
		void operator() (Args1 && ... args)
		{
			if (promise_.canceled())
				return;

			try
			{
				if constexpr (std::is_same_v<R, void>)
				{
					func_(std::forward<Args1>(args)...);
					promise_.set_value();
				}
				else
				{
					promise_.set_value(func_(std::forward<Args1>(args)...));
				}
			}
			catch(...)
			{
				promise_.set_exception(std::current_exception());
			}
		}

	private:
		promise<R> promise_;
		util::function<R(Args...)> func_;
	};

	template <typename T>
	template <typename F>
	auto future<T>::then(F && f)
	{
		if (!state_) throw empty_future_error{};

		std::lock_guard lock{state_->value_mutex};
		if (state_->value || state_->exception)
		{
			if constexpr (std::is_same_v<T, void>)
			{
				using R = decltype(f());
				promise<R> p;

				if (state_->value)
				{
					if constexpr (std::is_same_v<R, void>)
					{
						std::forward<F>(f)();
						p.set_value();
					}
					else
					{
						p.set_value(std::forward<F>(f)());
					}
				}
				else
					p.set_exception(state_->exception);
				return p.get_future();
			}
			else
			{
				using R = decltype(f(*(state_->value)));
				promise<R> p;

				if (state_->value)
				{
					if constexpr (std::is_same_v<R, void>)
					{
						std::forward<F>(f)(*(state_->value));
						p.set_value();
					}
					else
					{
						p.set_value(std::forward<F>(f)(*(state_->value)));
					}
				}
				else
					p.set_exception(state_->exception);
				return p.get_future();
			}
		}

		if constexpr (std::is_same_v<T, void>)
		{
			using R = decltype(f());
			packaged_task<R()> t(cancel_, std::forward<F>(f));
			auto fut = t.get_future();

			std::lock_guard lock{state_->then_mutex};
			state_->then_funcs.push_back(std::move(t));
			return fut;
		}
		else
		{
			using R = decltype(f(*(state_->value)));
			packaged_task<R(T &)> t(cancel_, std::forward<F>(f));
			auto fut = t.get_future();

			std::lock_guard lock{state_->then_mutex};
			state_->then_funcs.push_back(std::move(t));
			return fut;
		}
	}

}
