package de.arxsilex.maptools.converter.core.pbf

import java.io.ByteArrayOutputStream
import java.nio.file.Files
import java.util.zip.DeflaterOutputStream
import kotlin.io.path.createTempFile
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class PrimitiveBlockDecoderTest {
    @Test
    fun `decodes dense nodes tags coordinates and way references`() {
        val primitiveBlock = samplePrimitiveBlock()
        val nodes = mutableListOf<OsmPreviewNode>()
        val ways = mutableListOf<OsmPreviewWay>()
        val relations = mutableListOf<OsmPreviewRelation>()

        val statistics = PrimitiveBlockDecoder.decode(primitiveBlock, nodes::add, ways::add, relations::add)

        assertEquals(2, statistics.nodeCount)
        assertEquals(1, statistics.wayCount)
        assertEquals(100L, nodes[0].id)
        assertEquals(50.0, nodes[0].latitude)
        assertEquals(8.0, nodes[0].longitude)
        assertEquals("bench", nodes[0].tags["amenity"])
        assertEquals(105L, nodes[1].id)
        assertEquals(50.0001, nodes[1].latitude, 0.000000001)
        assertEquals(8.0002, nodes[1].longitude, 0.000000001)
        assertTrue(nodes[1].tags.isEmpty())
        assertEquals(200L, ways.single().id)
        assertTrue(longArrayOf(100, 105).contentEquals(ways.single().nodeIds))
        assertEquals("primary", ways.single().tags["highway"])
        assertEquals("Teststraße", ways.single().tags["name"])
        assertEquals(1, statistics.relationCount)
        assertEquals(300L, relations.single().id)
        assertEquals("route", relations.single().tags["type"])
        assertEquals(200L, relations.single().members.single().id)
        assertEquals(OsmRelationMemberType.WAY, relations.single().members.single().type)
        assertEquals("outer", relations.single().members.single().role)
    }

    @Test
    fun `collects a complete primitive block for shard processing`() {
        val data = PrimitiveBlockDecoder.decodeToData(samplePrimitiveBlock())

        assertEquals(2, data.nodes.size)
        assertEquals(1, data.ways.size)
        assertEquals(1, data.relations.size)
        assertEquals(PrimitiveBlockStatistics(2, 1, 1), data.statistics)
    }

    @Test
    fun `reads and inflates indexed data block from file`() {
        val primitiveBlock = samplePrimitiveBlock()
        val fileBytes = decoderPbfBlock("OSMHeader", byteArrayOf(1), false) +
            decoderPbfBlock("OSMData", primitiveBlock, true)
        val path = createTempFile(suffix = ".osm.pbf")
        Files.write(path, fileBytes)

        try {
            val index = PbfBlockScanner.scan(path)
            val unpacked = PbfDataBlockReader.read(path, index.entries[1])

            assertTrue(primitiveBlock.contentEquals(unpacked))
            assertEquals(2, PrimitiveBlockDecoder.decode(unpacked).nodeCount)
        } finally {
            Files.deleteIfExists(path)
        }
    }
}

internal fun samplePrimitiveBlock(): ByteArray {
    val strings = listOf("", "highway", "primary", "name", "Teststraße", "amenity", "bench", "type", "route", "outer")
    val stringTable = strings.fold(ByteArray(0)) { result, value ->
        result + decoderFieldBytes(1, value.toByteArray(Charsets.UTF_8))
    }
    val denseNodes = decoderFieldPackedSInt64(1, longArrayOf(100, 5)) +
        decoderFieldPackedSInt64(8, longArrayOf(500_000_000, 1_000)) +
        decoderFieldPackedSInt64(9, longArrayOf(80_000_000, 2_000)) +
        decoderFieldPackedUInt(10, intArrayOf(5, 6, 0, 0))
    val way = decoderFieldVarint(1, 200) +
        decoderFieldPackedUInt(2, intArrayOf(1, 3)) +
        decoderFieldPackedUInt(3, intArrayOf(2, 4)) +
        decoderFieldPackedSInt64(8, longArrayOf(100, 5))
    val relation = decoderFieldVarint(1, 300) +
        decoderFieldPackedUInt(2, intArrayOf(7)) +
        decoderFieldPackedUInt(3, intArrayOf(8)) +
        decoderFieldPackedUInt(8, intArrayOf(9)) +
        decoderFieldPackedSInt64(9, longArrayOf(200)) +
        decoderFieldPackedUInt(10, intArrayOf(1))
    val group = decoderFieldBytes(2, denseNodes) + decoderFieldBytes(3, way) + decoderFieldBytes(4, relation)
    return decoderFieldBytes(1, stringTable) +
        decoderFieldBytes(2, group) +
        decoderFieldVarint(17, 100)
}

internal fun decoderPbfBlock(type: String, unpacked: ByteArray, compressed: Boolean): ByteArray {
    val blob = if (compressed) {
        decoderFieldVarint(2, unpacked.size.toLong()) + decoderFieldBytes(3, unpacked.decoderZlib())
    } else {
        decoderFieldBytes(1, unpacked)
    }
    val header = decoderFieldBytes(1, type.toByteArray()) + decoderFieldVarint(3, blob.size.toLong())
    return byteArrayOf(
        (header.size ushr 24).toByte(),
        (header.size ushr 16).toByte(),
        (header.size ushr 8).toByte(),
        header.size.toByte(),
    ) + header + blob
}

private fun decoderFieldBytes(field: Int, bytes: ByteArray): ByteArray =
    decoderVarint((field shl 3 or 2).toLong()) + decoderVarint(bytes.size.toLong()) + bytes

private fun decoderFieldVarint(field: Int, value: Long): ByteArray =
    decoderVarint((field shl 3).toLong()) + decoderVarint(value)

private fun decoderFieldPackedUInt(field: Int, values: IntArray): ByteArray =
    decoderFieldBytes(field, values.fold(ByteArray(0)) { result, value -> result + decoderVarint(value.toLong()) })

private fun decoderFieldPackedSInt64(field: Int, values: LongArray): ByteArray =
    decoderFieldBytes(
        field,
        values.fold(ByteArray(0)) { result, value ->
            result + decoderVarint((value shl 1) xor (value shr 63))
        },
    )

private fun decoderVarint(value: Long): ByteArray {
    var remaining = value
    val output = ByteArrayOutputStream()
    while (remaining and -128L != 0L) {
        output.write(((remaining and 127L) or 128L).toInt())
        remaining = remaining ushr 7
    }
    output.write(remaining.toInt())
    return output.toByteArray()
}

private fun ByteArray.decoderZlib(): ByteArray {
    val output = ByteArrayOutputStream()
    DeflaterOutputStream(output).use { it.write(this) }
    return output.toByteArray()
}
