#include "ui.h"

UIGrid uiGridFromRoot(i32 width, i32 height) {
    return (UIGrid) {.bounds = {0, 0, width, height}};
}
UIGrid uiGridFromGrid(const UIGrid grid, i32 row, i32 col) {
    return uiGridFromGridSpan(grid, row, col, 1, 1);
}
UIGrid uiGridFromGridSpan(const UIGrid grid, i32 row, i32 col, i32 rowSpan, i32 colSpan) {
    Rectangle bounds = uiGridCellSpan(grid, row, col, rowSpan, colSpan);
    return (UIGrid) {
        .rowPadding = grid.rowPadding,
        .cellPadding = grid.cellPadding,
        .bounds = bounds
    };
}

Rectangle uiGridCell(const UIGrid grid, i32 row, i32 col) {
    return uiGridCellSpan(grid, row, col, 1, 1);
}
Rectangle uiGridCellSpan(const UIGrid grid, i32 row, i32 col, i32 rowSpan, i32 colSpan) {
    BZ_ASSERT(grid.rows > 0 && grid.cols > 0);
    BZ_ASSERT(rowSpan > 0 && colSpan > 0);
    BZ_ASSERT(row >= 0 && col >= 0);
    BZ_ASSERT(row + rowSpan <= grid.rows);
    BZ_ASSERT(col + colSpan <= grid.cols);

    f32 cellWidth = (grid.bounds.width - (grid.cols - 1) * grid.cellPadding) / grid.cols;
    f32 cellHeight = (grid.bounds.height - (grid.rows - 1) * grid.rowPadding) / grid.rows;

    Rectangle cellBounds = {
            .width  = cellWidth * colSpan,
            .height = cellHeight * rowSpan,
            .x      = (cellWidth + grid.cellPadding) * col,
            .y      = (cellHeight + grid.rowPadding) * row,
    };
    return cellBounds;
}

void uiGridDrawCells(const UIGrid grid) {
    BZ_ASSERT(grid.rows > 0 && grid.cols > 0);
    DrawRectangleLines(grid.bounds.x, grid.bounds.y, grid.bounds.width, grid.bounds.height, RED);

    for (i32 rowIdx = 0; rowIdx < grid.rows; rowIdx++) {
        for (i32 colIdx = 0; colIdx < grid.cols; colIdx++) {
            Rectangle rec = uiGridCell(grid, rowIdx, colIdx);
            DrawRectangleLines(rec.x, rec.y, rec.width, rec.height, BLUE);
        }
    }

}

static f32 getX(const Rectangle rec) {
    return rec.x;
}
static f32 getY(const Rectangle rec) {
    return rec.y;
}
static f32 getWidth(const Rectangle rec) {
    return rec.width;
}
static f32 getHeight(const Rectangle rec) {
    return rec.height;
}

typedef f32 (*FlexBoxGetFunc)(const Rectangle rec);

void uiFlexBoxCalc(const FlexBoxDesc *desc, i32 numElements, ...) {
    if (numElements <= 0) return;
    f32 totalWidth = 0;
    f32 totalHeight = 0;

    UIFlexBoxFlag flags = desc->flags;

    // Calculate total size of elements in the main axis
    va_list ap;
    va_start(ap, numElements);
    for (i32 i = 0; i < numElements; i++) {
        const Rectangle *rec = va_arg(ap, Rectangle *);
        totalWidth += rec->width;
        totalHeight += rec->height;
    }
    va_end(ap);

    // Default: DIR_ROW
    FlexBoxGetFunc axisX = getX, axisY = getY, axisWidth = getWidth, axisHeight = getHeight;
    FlexBoxGetFunc crossX = getY, crossY = getX, crossWidth = getHeight, crossHeight = getWidth;
    f32 mainAxisSize = totalWidth;
    if (flags & UI_FLEX_DIR_COLUMN) {
        axisX = getY;
        axisY = getX;
        axisWidth = getHeight;
        axisHeight = getWidth;
        crossX = getX;
        crossY = getY;
        crossWidth = getWidth;
        crossHeight = getHeight;
        mainAxisSize = totalHeight;
    }


    f32 spacingWidth = 0.0f;
    f32 spacingHeight = 0.0f;

    f32 mainAxisTotalSpacing = getWidth(desc->window) - mainAxisSize;
    f32 mainAxisPadding = desc->padding * (numElements - 1);

    f32 mainAxisSpacing = mainAxisTotalSpacing - mainAxisPadding;
    if (mainAxisSpacing < 0) mainAxisSpacing = 0;

    // Default: JUSTIFY_START
    f32 mainAxisOffset = 0;
    if (flags & UI_FLEX_JUSTIFY_CENTER) {
        mainAxisOffset = mainAxisSpacing * 0.5f;
    } else if (flags & UI_FLEX_JUSTIFY_END) {
        mainAxisOffset = mainAxisSpacing;
    }

    bool rowDir = flags & UI_FLEX_DIR_ROW;

    f32 offsetX = desc->window.x;
    f32 offsetY = desc->window.y;
    if (rowDir) offsetX += mainAxisOffset;
    else offsetY += mainAxisOffset;

    va_start(ap, numElements);
    for (i32 i = 0; i < numElements; i++) {
        f32 x = offsetX;
        f32 y = offsetY;

        Rectangle *rec = va_arg(ap, Rectangle *);
        f32 crossOffset = 0;
        if (flags & UI_FLEX_ALIGN_CENTER) {
            crossOffset = (getHeight(desc->window) - getHeight(*rec)) * 0.5f;
        } else if (flags & UI_FLEX_ALIGN_END) {
            crossOffset = (getHeight(desc->window) - getHeight(*rec));
        }
        if (rowDir) {
            offsetX += rec->width + spacingWidth + desc->padding;
            y += crossOffset;
        } else {
            offsetY += rec->height + spacingHeight + desc->padding;
            x += crossOffset;
        }


        rec->x = x;
        rec->y = y;

    }
    va_end(ap);
}