#include "ftxui/dom/box_helper.hpp"

#include <algorithm>  // for max

namespace ftxui::box_helper {

namespace {
// Called when the size allowed is greater than the requested size. This
// distributes the extra spaces toward the flexible elements, in relative
// proportions.
void ComputeGrow(std::vector<Element>* elements,
                 int extra_space,
                 int flex_grow_sum) {
  for (Element& element : *elements) {
    int added_space =
        extra_space * element.flex_grow / std::max(flex_grow_sum, 1);
    extra_space -= added_space;
    flex_grow_sum -= element.flex_grow;
    element.size = element.min_size + added_space;
  }
}

// Called when the size allowed is lower than the requested size, and the
// shrinkable element can absorbe the (negative) extra_space. This distribute
// the extra_space toward those.
void ComputeShrinkEasy(std::vector<Element>* elements,
                       int extra_space,
                       int flex_shrink_sum) {
  for (Element& element : *elements) {
    int added_space = extra_space * element.min_size * element.flex_shrink /
                      std::max(flex_shrink_sum, 1);
    extra_space -= added_space;
    flex_shrink_sum -= element.flex_shrink * element.min_size;
    element.size = element.min_size + added_space;
  }
}

// Called when the size allowed is lower than the requested size, and the
// shrinkable element can not absorbe the (negative) extra_space. This assign
// zero to shrinkable elements and distribute the remaining (negative)
// extra_space toward the other non shrinkable elements.
void ComputeShrinkHard(std::vector<Element>* elements,
                       int extra_space,
                       int size) {
  for (Element& element : *elements) {
    if (element.flex_shrink != 0) {
      element.size = 0;
      continue;
    }

    int added_space = extra_space * element.min_size / std::max(1, size);
    extra_space -= added_space;
    size -= element.min_size;

    element.size = element.min_size + added_space;
  }
}

}  // namespace

void Compute(std::vector<Element>* elements, int target_size) {
  int size = 0;
  int flex_grow_sum = 0;
  int flex_shrink_sum = 0;
  int flex_shrink_size = 0;

  for (auto& element : *elements) {
    flex_grow_sum += element.flex_grow;
    flex_shrink_sum += element.min_size * element.flex_shrink;
    if (element.flex_shrink != 0) {
      flex_shrink_size += element.min_size;
    }
    size += element.min_size;
  }

  int extra_space = target_size - size;
  if (extra_space >= 0) {
    ComputeGrow(elements, extra_space, flex_grow_sum);
  } else if (flex_shrink_size + extra_space >= 0) {
    ComputeShrinkEasy(elements, extra_space, flex_shrink_sum);

  } else {
    ComputeShrinkHard(elements, extra_space + flex_shrink_size,
                      size - flex_shrink_size);
  }
}

}  // namespace ftxui::box_helper

// Copyright 2021 Arthur Sonzogni. All rights reserved.
// Use of this source code is governed by the MIT license that can be found in
// the LICENSE file.