#include "StackLayout.h"
#include <algorithm>

namespace UI {

void StackLayout::Measure(Size availableSpace) {
    float availW = std::max(0.0f, availableSpace.width - margin.HorizontalTotal() - padding.HorizontalTotal());
    float availH = std::max(0.0f, availableSpace.height - margin.VerticalTotal() - padding.VerticalTotal());

    float accumPrimary = 0.0f;
    float maxSecondary = 0.0f;
    int visibleChildCount = 0;

    for (auto& child : children) {
        if (!child->IsVisible()) continue;
        visibleChildCount++;

        Size constraint{availW, availH};
        if (orientation == Orientation::Vertical) {
            constraint.height = std::max(0.0f, availH - accumPrimary);
        } else {
            constraint.width = std::max(0.0f, availW - accumPrimary);
        }

        child->Measure(constraint);
        Size desired = child->GetDesiredSize();

        if (orientation == Orientation::Vertical) {
            accumPrimary += desired.height;
            maxSecondary = std::max(maxSecondary, desired.width);
        } else {
            accumPrimary += desired.width;
            maxSecondary = std::max(maxSecondary, desired.height);
        }
    }

    if (visibleChildCount > 1) {
        accumPrimary += spacing * (visibleChildCount - 1);
    }

    if (orientation == Orientation::Vertical) {
        desiredSize.width = std::clamp(maxSecondary + padding.HorizontalTotal() + margin.HorizontalTotal(), minSize.width, maxSize.width);
        desiredSize.height = std::clamp(accumPrimary + padding.VerticalTotal() + margin.VerticalTotal(), minSize.height, maxSize.height);
    } else {
        desiredSize.width = std::clamp(accumPrimary + padding.HorizontalTotal() + margin.HorizontalTotal(), minSize.width, maxSize.width);
        desiredSize.height = std::clamp(maxSecondary + padding.VerticalTotal() + margin.VerticalTotal(), minSize.height, maxSize.height);
    }
}

void StackLayout::Arrange(Rect finalBounds) { // TODO: this method is way too long. needs refactor/antiantigravity
    ArrangeCore(finalBounds);

    Rect contentBounds{padding.Left, padding.Top, std::max(0.0f, localBounds.Width - padding.HorizontalTotal()), std::max(0.0f, localBounds.Height - padding.VerticalTotal())};

    float totalDesiredPrimary = 0.0f;
    int visibleChildCount = 0;

    struct StretchCandidate {
        Widget* widget;
        float baseSize;
        float maxSize;
        float allocatedSize;
        bool active;
    };
    std::vector<StretchCandidate> candidates;

    for (auto& child : children) {
        if (!child->IsVisible()) continue;
        visibleChildCount++;

        Size childDesired = child->GetDesiredSize();
        Thickness childMargin = child->GetMargin();

        if (orientation == Orientation::Vertical) {
            float baseH = std::max(0.0f, childDesired.height - childMargin.VerticalTotal());
            totalDesiredPrimary += childDesired.height;
            if (child->GetVerticalAlignment() == VerticalAlignment::Stretch && child->GetMaxSize().height > childDesired.height) {
                candidates.push_back({child.get(), baseH, child->GetMaxSize().height, baseH, true});
            }
        } else {
            float baseW = std::max(0.0f, childDesired.width - childMargin.HorizontalTotal());
            totalDesiredPrimary += childDesired.width;
            if (child->GetHorizontalAlignment() == HorizontalAlignment::Stretch && child->GetMaxSize().width > childDesired.width) {
                candidates.push_back({child.get(), baseW, child->GetMaxSize().width, baseW, true});
            }
        }
    }

    if (visibleChildCount > 1) {
        totalDesiredPrimary += spacing * (visibleChildCount - 1);
    }

    float availPrimary = (orientation == Orientation::Vertical) ? contentBounds.Height : contentBounds.Width;
    float remainingExcess = std::max(0.0f, availPrimary - totalDesiredPrimary);

    while (remainingExcess > 0.001f) {
        int activeCount = 0;
        for (auto& cand : candidates) {
            if (cand.active) activeCount++;
        }
        if (activeCount == 0) break;

        float share = remainingExcess / static_cast<float>(activeCount);
        bool anyCapped = false;

        for (auto& cand : candidates) {
            if (!cand.active) continue;
            float potential = cand.allocatedSize + share;
            if (potential >= cand.maxSize) {
                float actualAdd = cand.maxSize - cand.allocatedSize;
                cand.allocatedSize = cand.maxSize;
                remainingExcess -= actualAdd;
                cand.active = false;
                anyCapped = true;
            }
        }

        if (!anyCapped) {
            for (auto& cand : candidates) {
                if (cand.active) {
                    cand.allocatedSize += share;
                }
            }
            remainingExcess = 0.0f;
            break;
        }
    }

    auto getAllocatedSize = [&](Widget* w, float defaultSize) -> float { // TODO: make this not trash/review near ui-v2 complete
        for (auto& cand : candidates) {
            if (cand.widget == w) {
                return cand.allocatedSize;
            }
        }
        return defaultSize;
    };

    float currentOffset = 0.0f;

    for (auto& child : children) {
        if (!child->IsVisible()) continue;

        Size childDesired = child->GetDesiredSize();
        Thickness childMargin = child->GetMargin();

        float childX, childY, childW, childH;

        if (orientation == Orientation::Vertical) {
            float availCross = std::max(0.0f, contentBounds.Width - childMargin.HorizontalTotal());
            float naturalW   = std::max(0.0f, childDesired.width - childMargin.HorizontalTotal());

            auto horiz = CalculateAxisAlignment(child->GetHorizontalAlignment(), availCross, naturalW, childMargin.Left);
            childX = contentBounds.X + horiz.offset;
            childW = horiz.size;

            float baseH = std::max(0.0f, childDesired.height - childMargin.VerticalTotal());
            float finalH = getAllocatedSize(child.get(), baseH);
            finalH = std::clamp(finalH, child->GetMinSize().height, child->GetMaxSize().height);

            childY = contentBounds.Y + currentOffset + childMargin.Top;
            childH = finalH;

            currentOffset += childH + childMargin.VerticalTotal() + spacing;
        } else {
            float availCross = std::max(0.0f, contentBounds.Height - childMargin.VerticalTotal());
            float naturalH   = std::max(0.0f, childDesired.height - childMargin.VerticalTotal());

            auto vert = CalculateAxisAlignment(child->GetVerticalAlignment(), availCross, naturalH, childMargin.Top);
            childY = contentBounds.Y + vert.offset;
            childH = vert.size;

            float baseW = std::max(0.0f, childDesired.width - childMargin.HorizontalTotal());
            float finalW = getAllocatedSize(child.get(), baseW);
            finalW = std::clamp(finalW, child->GetMinSize().width, child->GetMaxSize().width);

            childX = contentBounds.X + currentOffset + childMargin.Left;
            childW = finalW;

            currentOffset += childW + childMargin.HorizontalTotal() + spacing;
        }

        child->Arrange(Rect(childX, childY, childW, childH));
    }

    ClearDirty(DirtyFlags::Layout);
}

}