package de.arxsilex.maptools.converter.desktop

import de.arxsilex.maptools.converter.core.GeoRectangleRegion
import de.arxsilex.maptools.converter.core.MapSelectionShape
import de.arxsilex.maptools.converter.core.toGeoRegions
import de.arxsilex.maptools.converter.core.unionAreaSquareKilometers
import de.arxsilex.maptools.converter.core.GeoBounds
import de.arxsilex.maptools.converter.core.preview.PreviewDataset
import kotlin.math.roundToLong

internal data class SelectionEstimate(
    val previewObjectCount: Int,
    val estimatedSourceObjectCount: Long,
    val estimatedSizeLowerBytes: Long,
    val estimatedSizeUpperBytes: Long,
)

internal data class SelectionEstimateSummary(
    val areas: List<SelectionEstimate>,
    val total: SelectionEstimate,
)

internal object PreviewSelectionEstimator {
    fun estimate(
        dataset: PreviewDataset,
        renderIndex: PreviewRenderIndex,
        shapes: List<MapSelectionShape>,
        sourceBounds: GeoBounds,
        sourceSizeBytes: Long,
    ): SelectionEstimateSummary {
        if (shapes.isEmpty()) {
            return SelectionEstimateSummary(emptyList(), SelectionEstimate(0, 0, 0, 0))
        }
        val totalPreviewObjects = dataset.points.size + dataset.polylines.size
        val sourceObjects = dataset.sourceNodeCount + dataset.sourceWayCount
        val unionAreas = HashSet<Int>()
        val unionLines = HashSet<Int>()
        val unionPoints = HashSet<Int>()

        fun indices(shape: MapSelectionShape): SelectedPreviewIndices {
            val area = shape.normalizedEnvelope(sourceBounds)
            val region = shape.toGeoRegion(sourceBounds)
            fun polylineInside(index: Int): Boolean {
                val coordinates = dataset.polylines[index].coordinates
                var coordinateIndex = 0
                while (coordinateIndex < coordinates.size) {
                    if (region.contains(coordinates[coordinateIndex + 1], coordinates[coordinateIndex])) return true
                    coordinateIndex += 2
                }
                return false
            }
            return SelectedPreviewIndices(
                areaIndices = renderIndex.queryAreaIndices(area.left, area.top, area.right, area.bottom)
                    .filter(::polylineInside).toSet(),
                lineIndices = renderIndex.queryLineIndices(area.left, area.top, area.right, area.bottom)
                    .filter(::polylineInside).toSet(),
                pointIndices = renderIndex.queryPointIndices(area.left, area.top, area.right, area.bottom)
                    .filter { index -> dataset.points[index].let { region.contains(it.latitude, it.longitude) } }
                    .toSet(),
            )
        }

        val areaEstimates = shapes.map { shape ->
            val selected = indices(shape)
            unionAreas += selected.areaIndices
            unionLines += selected.lineIndices
            unionPoints += selected.pointIndices
            estimateFor(
                previewObjects = selected.count,
                totalPreviewObjects = totalPreviewObjects,
                fallbackAreaFraction = listOf(shape.toGeoRegion(sourceBounds)).unionAreaSquareKilometers(sourceBounds) /
                    listOf(GeoRectangleRegion(sourceBounds)).unionAreaSquareKilometers(),
                sourceObjects = sourceObjects,
                sourceSizeBytes = sourceSizeBytes,
            )
        }
        val unionPreviewCount = unionAreas.size + unionLines.size + unionPoints.size
        return SelectionEstimateSummary(
            areas = areaEstimates,
            total = estimateFor(
                previewObjects = unionPreviewCount,
                totalPreviewObjects = totalPreviewObjects,
                fallbackAreaFraction = shapes.toGeoRegions(sourceBounds).unionAreaSquareKilometers(sourceBounds) /
                    listOf(GeoRectangleRegion(sourceBounds)).unionAreaSquareKilometers(),
                sourceObjects = sourceObjects,
                sourceSizeBytes = sourceSizeBytes,
            ),
        )
    }

    private fun estimateFor(
        previewObjects: Int,
        totalPreviewObjects: Int,
        fallbackAreaFraction: Double,
        sourceObjects: Long,
        sourceSizeBytes: Long,
    ): SelectionEstimate {
        val ratio = if (totalPreviewObjects > 0) {
            previewObjects.toDouble() / totalPreviewObjects
        } else {
            fallbackAreaFraction
        }.coerceIn(0.0, 1.0)
        val expectedBytes = sourceSizeBytes * ratio
        return SelectionEstimate(
            previewObjectCount = previewObjects,
            estimatedSourceObjectCount = (sourceObjects * ratio).roundToLong(),
            estimatedSizeLowerBytes = (expectedBytes * 0.65).roundToLong(),
            estimatedSizeUpperBytes = (expectedBytes * 1.40).roundToLong(),
        )
    }

    private data class SelectedPreviewIndices(
        val areaIndices: Set<Int>,
        val lineIndices: Set<Int>,
        val pointIndices: Set<Int>,
    ) {
        val count: Int get() = areaIndices.size + lineIndices.size + pointIndices.size
    }
}
