#include "VirtualizingGridLayout.h"
#include <algorithm>
#include <cmath>

namespace UI {

VirtualizingGridLayout::VirtualizingGridLayout(
    std::size_t count,
    float width,
    float itemSpacing,
    float rowHeight,
    float itemOverscan,
    std::function<Size(std::size_t)> sizeProvider,
    std::function<std::unique_ptr<Widget>(std::size_t)> factory
) : itemCount(count), reasonableColumnWidth(std::max(1.0f, width)), spacing(std::max(0.0f, itemSpacing)),
    maximumRowHeight(std::max(0.0f, rowHeight)),
    overscan(std::max(0.0f, itemOverscan)), itemSizeProvider(std::move(sizeProvider)), itemFactory(std::move(factory)) {
}

std::size_t VirtualizingGridLayout::CalculateColumnCount(float availableWidth) const {
    float contentWidth = std::max(0.0f, availableWidth - padding.HorizontalTotal());
    if (contentWidth <= 0.0f) {
        return 1;
    }
    float columns = (contentWidth + spacing) / (reasonableColumnWidth + spacing);
    return std::max<std::size_t>(1, static_cast<std::size_t>(std::round(columns)));
}

float VirtualizingGridLayout::CalculateColumnWidth(float availableWidth) const {
    float contentWidth = std::max(0.0f, availableWidth - padding.HorizontalTotal());
    return std::max(0.0f, (contentWidth - spacing * static_cast<float>(columnCount - 1)) /
        static_cast<float>(columnCount));
}

void VirtualizingGridLayout::RebuildMetrics(float availableWidth) {
    std::size_t previousColumnCount = columnCount;
    columnCount = CalculateColumnCount(availableWidth);
    if (columnCount != previousColumnCount) {
        hasVisibleRange = false;
    }
    rowHeights.clear();
    rowOffsets.clear();
    std::size_t rowCount = itemCount == 0 ? 0 : (itemCount + columnCount - 1) / columnCount;
    rowHeights.resize(rowCount, 0.0f);
    rowOffsets.resize(rowCount, padding.Top);
    float rowHeight = 0.0f;
    for (std::size_t index = 0; index < itemCount; index++) {
        Size itemSize = itemSizeProvider ? itemSizeProvider(index) : Size();
        rowHeight = std::max(rowHeight, std::max(0.0f, itemSize.height));
    }
    if (maximumRowHeight > 0.0f) {
        rowHeight = maximumRowHeight;
    }
    for (std::size_t row = 0; row < rowCount; row++) {
        rowHeights[row] = rowHeight;
    }
    for (std::size_t row = 1; row < rowCount; row++) {
        rowOffsets[row] = rowOffsets[row - 1] + rowHeights[row - 1] + spacing;
    }
}

bool VirtualizingGridLayout::UpdateVisibleChildren() {
    if (!itemCount || !itemFactory || rowHeights.empty() || !parent) {
        children.clear();
        hasVisibleRange = false;
        return true;
    }
    float viewportHeight = std::max(0.0f, parent->GetLocalBounds().Height - parent->GetPadding().VerticalTotal());
    float start = std::max(0.0f, scrollOffset - overscan);
    float end = scrollOffset + viewportHeight + overscan;
    auto firstIterator = std::upper_bound(rowOffsets.begin(), rowOffsets.end(), start);
    std::size_t first = firstIterator == rowOffsets.begin() ? 0 :
        static_cast<std::size_t>(std::distance(rowOffsets.begin(), firstIterator) - 1);
    auto lastIterator = std::upper_bound(rowOffsets.begin(), rowOffsets.end(), end);
    std::size_t last = lastIterator == rowOffsets.begin() ? 0 :
        static_cast<std::size_t>(std::distance(rowOffsets.begin(), lastIterator) - 1);
    last = std::min(last, rowHeights.size() - 1);
    if (hasVisibleRange && first == visibleStartRow && last == visibleEndRow) {
        return false;
    }
    children.clear();
    visibleStartRow = first;
    visibleEndRow = last;
    hasVisibleRange = true;
    for (std::size_t row = first; row <= last; row++) {
        std::size_t firstIndex = row * columnCount;
        std::size_t lastIndex = std::min(itemCount, firstIndex + columnCount);
        for (std::size_t index = firstIndex; index < lastIndex; index++) {
            std::unique_ptr<Widget> child = itemFactory(index);
            if (child) {
                AddChild(std::move(child));
            }
        }
    }
    return true;
}

void VirtualizingGridLayout::Measure(Size availableSpace) {
    RebuildMetrics(availableSpace.width);
    float totalHeight = padding.VerticalTotal();
    if (!rowHeights.empty()) {
        totalHeight += rowOffsets.back() - padding.Top + rowHeights.back();
    }
    desiredSize.width = std::clamp(availableSpace.width + margin.HorizontalTotal(), minSize.width, maxSize.width);
    desiredSize.height = std::clamp(totalHeight + margin.VerticalTotal(), minSize.height, maxSize.height);
}

void VirtualizingGridLayout::ArrangeChildInCell(Widget* child, Rect cellBounds) const {
    Size desired = child->GetDesiredSize();
    float scale = 1.0f;
    if (desired.width > 0.0f && desired.height > 0.0f) {
        scale = std::min(cellBounds.Width / desired.width, cellBounds.Height / desired.height);
    }
    float width = desired.width * scale;
    float height = desired.height * scale;
    child->Arrange(Rect(
        cellBounds.X + (cellBounds.Width - width) * 0.5f,
        cellBounds.Y + (cellBounds.Height - height) * 0.5f,
        width,
        height
    ));
}

void VirtualizingGridLayout::Arrange(Rect finalBounds) {
    ArrangeCore(finalBounds);
    RebuildMetrics(localBounds.Width);
    UpdateVisibleChildren();
    if (!hasVisibleRange) {
        ClearDirty(DirtyFlags::Layout | DirtyFlags::Hierarchy);
        return;
    }
    float columnWidth = CalculateColumnWidth(localBounds.Width);
    for (std::size_t childIndex = 0; childIndex < children.size(); childIndex++) {
        std::size_t itemIndex = visibleStartRow * columnCount + childIndex;
        if (itemIndex >= itemCount) {
            break;
        }
        std::size_t row = itemIndex / columnCount;
        std::size_t column = itemIndex % columnCount;
        Widget* child = children[childIndex].get();
        child->Measure(Size(columnWidth, rowHeights[row]));
        ArrangeChildInCell(child, Rect(
            padding.Left + static_cast<float>(column) * (columnWidth + spacing),
            rowOffsets[row], columnWidth, rowHeights[row]
        ));
    }
    ClearDirty(DirtyFlags::Layout | DirtyFlags::Hierarchy);
}

void VirtualizingGridLayout::SetScrollOffset(Point offset) {
    float newOffset = std::max(0.0f, offset.y);
    if (newOffset == scrollOffset) {
        return;
    }
    scrollOffset = newOffset;
    if (UpdateVisibleChildren() && localBounds.Width > 0.0f) {
        Arrange(localBounds);
    }
    MarkDirty(DirtyFlags::Render);
}

}