#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::weight{};

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

			float total_weight = 0.f;

			// First, allocate minimized & fixed elements
			// Also compute the total weight for weighted elements
			for (std::size_t i = 0; i < policies.size(); ++i)
			{
				auto const policy = policies[i] ? *policies[i] : default_policy;

				if (std::get_if<box_layout::minimized>(&policy))
				{
					result[i] = children[i]->size_constraints(Dimension, other_dimension_size).min;
					total_size -= result[i];
				}
				else if (auto fixed = std::get_if<box_layout::fixed>(&policy))
				{
					result[i] = fixed->value;
					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 < policies.size(); ++i)
			{
				auto const policy = policies[i] ? *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(react::value<float> margin)
		: margin_(margin)
	{}

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

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

		float const margin = *margin_;

		auto sizes = allocate<Dimension>(new_shape[Dimension].length() - (children().size() - 1) * margin, new_shape[Dimension ^ 1].length(), *size_policies_, children());

		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 += margin;
		}
	}

	template <int Dimension>
	geom::interval<float> box_layout_base<Dimension>::size_constraints(int dimension, float other_dimension_size) const
	{
		if (children().empty())
			return container::size_constraints(dimension, other_dimension_size);

		float const margin = *margin_;

		if (dimension == Dimension)
		{
			auto result = geom::interval{0.f, 0.f};

			auto const & policies = *size_policies_;

			auto weighted_unit_constraints = geom::interval{0.f, infinity};
			float weight_sum = 0.f;

			for (std::size_t i = 0; i < children().size(); ++i)
			{
				auto const & child = children()[i];
				auto const policy = policies[i] ? *policies[i] : default_policy;

				if (auto fixed = std::get_if<box_layout::fixed>(&policy))
					result += fixed->value;
				else if (std::get_if<box_layout::minimized>(&policy))
				{
					if (child)
					{
						auto child_constraints = child->size_constraints(dimension, other_dimension_size);
						result.min += child_constraints.min;
						result.max += child_constraints.max;
					}
					else
					{
						result.max = infinity;
					}
				}
				else if (auto weight = std::get_if<box_layout::weight>(&policy))
				{
					if (child)
					{
						auto child_constraints = child->size_constraints(dimension, other_dimension_size);
						child_constraints.min /= weight->value;
						child_constraints.max /= weight->value;
						weighted_unit_constraints &= child_constraints;
					}

					weight_sum += weight->value;
				}
			}

			// Prevent NaN in result.max (inf * 0)
			if (weight_sum > 0.f)
			{
				result.min += weighted_unit_constraints.min * weight_sum;
				result.max += weighted_unit_constraints.max * weight_sum;
			}

			result += (children().size() - 1) * margin;

			return result;
		}
		else
		{
			float const margin = *margin_;

			auto sizes = allocate<Dimension>(other_dimension_size - (children().size() - 1) * margin, shape()[1].length(), *size_policies_, children());

			auto result = container::size_constraints(dimension, other_dimension_size);

			for (std::size_t i = 0; i < children().size(); ++i)
			{
				if (!children()[i]) continue;
				result &= children()[i]->size_constraints(dimension, sizes[i]);
			}

			return result;
		}
	}

	template <int Dimension>
	void box_layout_base<Dimension>::update(typename detail::box_layout_type<Dimension>::type const & box_layout)
	{
		size_policies_ = 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);
	}

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

}