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.assertFailsWith
import kotlin.test.assertTrue

class PbfBlockScannerTest {
    @Test
    fun `indexes header and multiple data blocks`() {
        val header = pbfBlock("OSMHeader", byteArrayOf(1, 2, 3), compressed = false)
        val dataOne = pbfBlock("OSMData", ByteArray(120) { it.toByte() }, compressed = true)
        val dataTwo = pbfBlock("OSMData", ByteArray(80) { (it * 2).toByte() }, compressed = true)
        val fileBytes = header + dataOne + dataTwo
        val path = createTempFile(suffix = ".osm.pbf")
        Files.write(path, fileBytes)
        val progress = mutableListOf<PbfScanProgress>()

        try {
            val index = PbfBlockScanner.scan(path, onProgress = progress::add)

            assertEquals(fileBytes.size.toLong(), index.sourceSize)
            assertEquals(3, index.entries.size)
            assertEquals(2, index.dataBlockCount)
            assertEquals("OSMHeader", index.entries[0].type)
            assertEquals(PbfCompression.RAW, index.entries[0].compression)
            assertEquals(PbfCompression.ZLIB, index.entries[1].compression)
            assertEquals(120, index.entries[1].rawSize)
            assertEquals(index.entries[0].totalStoredSize(), index.entries[1].fileOffset)
            assertEquals(fileBytes.size.toLong(), progress.last().bytesRead)
            assertEquals(64, index.sourceSha256.length)
        } finally {
            Files.deleteIfExists(path)
        }
    }

    @Test
    fun `honors cancellation between blocks`() {
        val path = createTempFile(suffix = ".osm.pbf")
        Files.write(
            path,
            pbfBlock("OSMHeader", byteArrayOf(1), compressed = false) +
                pbfBlock("OSMData", ByteArray(30), compressed = true),
        )
        var callbacks = 0

        try {
            assertFailsWith<PbfScanCancelledException> {
                PbfBlockScanner.scan(
                    path = path,
                    onProgress = { callbacks++ },
                    isCancelled = { callbacks == 1 },
                )
            }
            assertEquals(1, callbacks)
        } finally {
            Files.deleteIfExists(path)
        }
    }

    @Test
    fun `rejects truncated blob`() {
        val valid = pbfBlock("OSMHeader", byteArrayOf(1, 2, 3), compressed = false)
        val path = createTempFile(suffix = ".osm.pbf")
        Files.write(path, valid.copyOf(valid.size - 1))

        try {
            val error = assertFailsWith<Exception> { PbfBlockScanner.scan(path) }
            assertTrue(error.message.orEmpty().contains("unerwartet"))
        } finally {
            Files.deleteIfExists(path)
        }
    }
}

private fun PbfBlockIndexEntry.totalStoredSize(): Long = 4L + blobHeaderSize + blobSize

private fun pbfBlock(type: String, unpacked: ByteArray, compressed: Boolean): ByteArray {
    val blob = if (compressed) {
        scannerFieldVarint(2, unpacked.size.toLong()) + scannerFieldBytes(3, unpacked.zlibForScanner())
    } else {
        scannerFieldBytes(1, unpacked)
    }
    val blobHeader = scannerFieldString(1, type) + scannerFieldVarint(3, blob.size.toLong())
    return byteArrayOf(
        (blobHeader.size ushr 24).toByte(),
        (blobHeader.size ushr 16).toByte(),
        (blobHeader.size ushr 8).toByte(),
        blobHeader.size.toByte(),
    ) + blobHeader + blob
}

private fun scannerFieldString(field: Int, value: String): ByteArray =
    scannerFieldBytes(field, value.toByteArray())

private fun scannerFieldBytes(field: Int, value: ByteArray): ByteArray =
    scannerVarint((field shl 3 or 2).toLong()) + scannerVarint(value.size.toLong()) + value

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

private fun scannerVarint(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.zlibForScanner(): ByteArray {
    val output = ByteArrayOutputStream()
    DeflaterOutputStream(output).use { it.write(this) }
    return output.toByteArray()
}
