package org.chsrobotics.dash

import androidx.compose.runtime.Composable
import androidx.compose.runtime.Stable
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.layout.Measurable
import androidx.compose.ui.layout.ParentDataModifier
import androidx.compose.ui.unit.Constraints
import androidx.compose.ui.unit.Density

data class GridLayout(
    val column: Int,
    val row: Int,
    val columnSpan: Int,
    val rowSpan: Int
)

interface GridScope {
    @Stable
    fun Modifier.grid(column: Int, row: Int, columnSpan: Int = 1, rowSpan: Int = 1) = this.then(
        GridData(column, row, columnSpan, rowSpan)
    )
    @Stable
    fun Modifier.grid(layout: GridLayout) = this.then(
        GridData(layout.column, layout.row, layout.columnSpan, layout.rowSpan)
    )

    companion object : GridScope
}

private class GridData(
    val column: Int,
    val row: Int,
    val columnSpan: Int,
    val rowSpan: Int,
) : ParentDataModifier {

    override fun Density.modifyParentData(parentData: Any?): Any = this@GridData

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (javaClass != other?.javaClass) return false

        other as GridData

        if (columnSpan != other.columnSpan) return false
        if (rowSpan != other.rowSpan) return false

        return true
    }

    override fun hashCode(): Int {
        var result = columnSpan
        result = 31 * result + rowSpan
        return result
    }
}

private val Measurable.gridData: GridData?
    get() = parentData as? GridData

private val Measurable.columnSpan: Int
    get() = gridData?.columnSpan ?: 1

private val Measurable.rowSpan: Int
    get() = gridData?.rowSpan ?: 1

data class GridInfo(
    val numChildren: Int,
    val columnSpan: Int,
    val rowSpan: Int,
)

@Composable
fun Grid(
    columns: Int,
    rows: Int,
    modifier: Modifier = Modifier,
    content: @Composable GridScope.() -> Unit,
) {
    println("CALLING GRId")
    check(columns > 0) { "Columns must be greater than 0" }
    check(rows > 0) { "Rows must be greater than 0" }
    Layout(
        content = { GridScope.content() },
        modifier = modifier,
    ) { measurables, constraints ->
        val standardGrid = GridData(0, 0, 1, 1)
        val spans = measurables.map { measurable -> measurable.gridData ?: standardGrid }

        spans.forEach {
            println(it.column)
        }

        // build constraints
        val baseConstraints = Constraints.fixed(
            width = constraints.maxWidth / columns,
            height = constraints.maxHeight / rows,
        )
        val cellConstraints = measurables.map { measurable ->
            val columnSpan = measurable.columnSpan
            val rowSpan = measurable.rowSpan
            Constraints.fixed(
                width = baseConstraints.maxWidth * columnSpan,
                height = baseConstraints.maxHeight * rowSpan
            )
        }

        // measure children
        val placeables = measurables.mapIndexed { index, measurable ->
            measurable.measure(cellConstraints[index])
        }

        layout(
            width = constraints.maxWidth,
            height = constraints.maxHeight,
        ) {
            placeables.forEachIndexed { i, placeable ->
                placeable.placeRelative(
                    x = spans[i].column * baseConstraints.maxWidth,
                    y = spans[i].row * baseConstraints.maxHeight
                )
            }
        }
    }
}