#include <psemek/ui/impl/box_layout_base.hpp>
#include <psemek/react/map.hpp>
#include <psemek/react/join.hpp>

namespace psemek::ui::impl
{

	namespace
	{

		static constexpr box_layout::size_policy default_policy = box_layout::minimized{};

		template <int Dimension>
		size_constraints compute_size_constraints(std::vector<box_layout::size_policy> const & size_policies,
			std::vector<size_constraints> const & children_size_constraints, float padding)
		{
			size_constraints minimized = size_constraints::max();
			size_constraints weight_unit = minimized;
			float weight_sum = 0.f;

			for (std::size_t i = 0; i < std::min(size_policies.size(), children_size_constraints.size()); ++i)
			{
				if (std::get_if<box_layout::minimized>(&size_policies[i]))
				{
					minimized = sum_along(minimized, children_size_constraints[i], Dimension);
				}
				else if (auto weight = std::get_if<box_layout::weight>(&size_policies[i]))
				{
					weight_unit = intersect(weight_unit, scale(children_size_constraints[i], 1.f / weight->value, Dimension));
					weight_sum += weight->value;
				}
			}

			auto result = std::move(minimized);

			if (weight_sum > 0.f)
				result = sum_along(result, scale(weight_unit, weight_sum, Dimension), Dimension);

			if (size_policies.size() > 0)
			{
				geom::vector shift_delta{0.f, 0.f};
				shift_delta[Dimension] = padding * (size_policies.size() - 1);
				result = shift(std::move(result), shift_delta);
			}

			return result;
		}

		template <int Dimension>
		std::vector<float> allocate(float total_size, std::vector<react::value<box_layout::size_policy>> const & size_policies,
			util::span<std::unique_ptr<component> const> children, float padding)
		{
			std::vector<float> result(std::min(size_policies.size(), children.size()), 0.f);

			if (!size_policies.empty())
				total_size -= padding * (result.size() - 1);

			float total_weight = 0.f;

			// First, allocate minimized elements
			// Also compute the total weight for weighted elements
			for (std::size_t i = 0; i < result.size(); ++i)
			{
				if (!children[i])
					continue;

				auto const policy = size_policies[i] ? *size_policies[i] : default_policy;

				if (std::get_if<box_layout::minimized>(&policy))
				{
					result[i] = min(*children[i]->size_constraints(), Dimension);
					total_size -= result[i];
				}
				else if (auto weight = std::get_if<box_layout::weight>(&policy))
				{
					total_weight += weight->value;
				}
			}

			// Next, allocate weighted elements
			for (std::size_t i = 0; i < result.size(); ++i)
			{
				auto const policy = size_policies[i] ? *size_policies[i] : default_policy;

				if (auto weight = std::get_if<box_layout::weight>(&policy))
				{
					result[i] = total_size * weight->value / total_weight;
				}
			}

			return result;
		}

	}

	template <int Dimension>
	box_layout_base<Dimension>::box_layout_base()
		: padding_(0.f)
		, size_policies_(std::vector<react::value<box_layout::size_policy>>{})
		, children_size_constraints_()
		, size_constraints_(
			react::map(
				[](auto const & size_policies, auto const & children_size_constraints, auto const & padding){
					return compute_size_constraints<Dimension>(size_policies, children_size_constraints, padding);
				},
				react::join(react::map(react::unpack_with_default(default_policy), react::join(size_policies_))),
				react::join(react::map(react::unpack, children_size_constraints_)),
				react::join(padding_)
			)
		)
	{}

	template <int Dimension>
	void box_layout_base<Dimension>::reshape(geom::box<float, 2> const & new_shape)
	{
		container::reshape(new_shape);

		if (children().empty())
			return;

		auto padding = **padding_;

		auto sizes = allocate<Dimension>(new_shape[Dimension].length(), **size_policies_, children(), padding);

		float pen = new_shape[Dimension].min;
		for (std::size_t i = 0; i < children().size(); ++i)
		{
			if (!children()[i]) continue;

			geom::box<float, 2> child_shape;
			child_shape[Dimension] = {pen, pen + sizes[i]};
			child_shape[Dimension ^ 1] = new_shape[Dimension ^ 1];
			children()[i]->reshape(child_shape);

			pen += sizes[i];
			pen += padding;
		}
	}

	template <int Dimension>
	react::value<size_constraints> box_layout_base<Dimension>::size_constraints() const
	{
		return size_constraints_;
	}

	template <int Dimension>
	void box_layout_base<Dimension>::set_children(std::vector<std::unique_ptr<component>> children)
	{
		container::set_children(std::move(children));

		std::vector<react::value<struct size_constraints>> children_size_constraints;
		for (auto const & child : this->children())
		{
			if (child)
				children_size_constraints.push_back(child->size_constraints());
			else
				children_size_constraints.push_back(size_constraints::max());
		}
		children_size_constraints_.set(std::move(children_size_constraints));
	}

	template <int Dimension>
	std::vector<std::unique_ptr<component>> box_layout_base<Dimension>::release_children()
	{
		auto result = container::release_children();
		children_size_constraints_.set({});
		return result;
	}

	template <int Dimension>
	void box_layout_base<Dimension>::update(typename detail::box_layout_type<Dimension>::type const & box_layout)
	{
		size_policies_.set(react::map([](auto const & children){
			std::vector<react::value<box_layout::size_policy>> size_policies;
			size_policies.reserve(children.size());
			for (auto const & child : children)
				size_policies.push_back(child.policy);
			return size_policies;
		}, box_layout.children));

		padding_.set(box_layout.padding);
	}

	template struct box_layout_base<0>;
	template struct box_layout_base<1>;

}