package space.nstart.pcp.slots_service.util

import org.locationtech.jts.geom.Coordinate
import org.locationtech.jts.geom.CoordinateSequenceFilter
import org.locationtech.jts.geom.Envelope
import org.locationtech.jts.geom.Geometry
import org.locationtech.jts.geom.GeometryCollection
import org.locationtech.jts.geom.GeometryFactory
import org.locationtech.jts.geom.LinearRing
import org.locationtech.jts.geom.MultiPolygon
import org.locationtech.jts.geom.Polygon
import space.nstart.pcp.pcp_types_lib.utils.wkt.WKTParser
import kotlin.math.abs
import kotlin.math.floor

object LongitudeWrapGeometry {

    private val parserThreadLocal = ThreadLocal.withInitial { WKTParser() }
    private val geometryFactory = GeometryFactory()

    private fun parser(): WKTParser = parserThreadLocal.get()

    fun normalizeToContinuous360Wkt(wkt: String): String =
        parser().toWKT(normalizeToContinuous360(parser().parseWKT(wkt)))

    fun buildSearchVariantsWkt(wkt: String): List<String> =
        buildSearchVariants(parser().parseWKT(wkt)).map(parser()::toWKT).distinct()

    fun normalizeToContinuous360(geometry: Geometry): Geometry {
        val normalized = shiftNegativeLongitudesTo360(geometry.copy())
        return when (normalized) {
            is Polygon -> normalizePolygon(normalized)
            is MultiPolygon -> normalizeMultiPolygon(normalized)
            else -> normalized
        }.let { if (it.isValid) it else it.buffer(0.0) }
    }

    fun buildSearchVariants(geometry: Geometry): List<Geometry> {
        val normalized = normalizeToContinuous360(geometry)
        return listOf(
            shiftLongitude(normalized, -360.0),
            normalized,
            shiftLongitude(normalized, 360.0)
        ).distinctBy { parser().toWKT(it) }
    }

    fun alignToReference(geometry: Geometry, reference: Geometry): Geometry {
        val normalizedReference = normalizeToContinuous360(reference)
        val referenceCenter = envelopeCenterX(normalizedReference.envelopeInternal)

        return buildSearchVariants(geometry)
            .minByOrNull { abs(envelopeCenterX(it.envelopeInternal) - referenceCenter) }
            ?: normalizeToContinuous360(geometry)
    }

    private fun shiftNegativeLongitudesTo360(geometry: Geometry): Geometry {
        geometry.apply(object : CoordinateSequenceFilter {
            override fun filter(seq: org.locationtech.jts.geom.CoordinateSequence, i: Int) {
                val lon = seq.getX(i)
                if (lon < 0.0) {
                    seq.setOrdinate(i, 0, lon + 360.0)
                }
            }

            override fun isDone(): Boolean = false

            override fun isGeometryChanged(): Boolean = true
        })
        geometry.geometryChanged()
        return geometry
    }

    private fun normalizePolygon(polygon: Polygon): Geometry {
        val shell = unwrapRing(polygon.exteriorRing)
        val holes = Array(polygon.numInteriorRing) { index ->
            unwrapRing(polygon.getInteriorRingN(index))
        }

        return shiftIntoCanonicalBand(repairGeometry(
            polygon.factory.createPolygon(shell, holes)
        ))
    }

    private fun normalizeMultiPolygon(multiPolygon: MultiPolygon): Geometry {
        val parts = mutableListOf<Polygon>()
        for (index in 0 until multiPolygon.numGeometries) {
            collectPolygons(normalizePolygon(multiPolygon.getGeometryN(index) as Polygon), parts)
        }

        return when (parts.size) {
            0 -> geometryFactory.createPolygon()
            1 -> parts.first()
            else -> geometryFactory.createMultiPolygon(parts.toTypedArray())
        }
    }

    private fun collectPolygons(geometry: Geometry, parts: MutableList<Polygon>) {
        when (geometry) {
            is Polygon -> parts += geometry
            is MultiPolygon -> {
                for (index in 0 until geometry.numGeometries) {
                    parts += geometry.getGeometryN(index) as Polygon
                }
            }
            is GeometryCollection -> {
                for (index in 0 until geometry.numGeometries) {
                    collectPolygons(geometry.getGeometryN(index), parts)
                }
            }
        }
    }

    private fun repairGeometry(geometry: Geometry): Geometry =
        if (geometry.isValid) geometry else geometry.buffer(0.0)

    private fun shiftIntoCanonicalBand(geometry: Geometry): Geometry {
        val minX = geometry.coordinates.minOfOrNull { it.x } ?: return geometry
        if (minX >= 0.0) return geometry

        val shift = 360.0 * (floor(-minX / 360.0) + 1.0)
        return shiftLongitude(geometry, shift)
    }

    private fun shiftLongitude(geometry: Geometry, delta: Double): Geometry {
        if (delta == 0.0) return geometry
        val shifted = geometry.copy()
        shifted.apply(object : CoordinateSequenceFilter {
            override fun filter(seq: org.locationtech.jts.geom.CoordinateSequence, i: Int) {
                seq.setOrdinate(i, 0, seq.getX(i) + delta)
            }

            override fun isDone(): Boolean = false

            override fun isGeometryChanged(): Boolean = true
        })
        shifted.geometryChanged()
        return shifted
    }

    private fun envelopeCenterX(envelope: Envelope): Double = (envelope.minX + envelope.maxX) / 2.0

    private fun unwrapRing(ring: org.locationtech.jts.geom.LineString): LinearRing {
        val source = ring.coordinates
        if (source.isEmpty()) {
            return ring.factory.createLinearRing()
        }

        val first = source.first()
        val unwrapped = ArrayList<Coordinate>(source.size)
        var previousX = first.x
        unwrapped.add(Coordinate(previousX, first.y, first.z))

        for (index in 1 until source.size - 1) {
            val current = source[index]
            var adjustedX = current.x

            while (adjustedX - previousX > 180.0) {
                adjustedX -= 360.0
            }
            while (adjustedX - previousX < -180.0) {
                adjustedX += 360.0
            }

            unwrapped.add(Coordinate(adjustedX, current.y, current.z))
            previousX = adjustedX
        }

        val start = unwrapped.first()
        unwrapped.add(Coordinate(start.x, start.y, start.z))
        return ring.factory.createLinearRing(unwrapped.toTypedArray())
    }
}
