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

import java.io.EOFException
import java.io.IOException
import java.io.InputStream
import java.nio.file.Files
import java.nio.file.Path
import java.security.DigestInputStream
import java.security.MessageDigest
import kotlin.io.path.fileSize
import kotlin.io.path.getLastModifiedTime

enum class PbfCompression {
    RAW,
    ZLIB,
    LZMA,
    BZIP2,
    LZ4,
    ZSTD,
}

data class PbfBlockIndexEntry(
    val sequence: Int,
    val fileOffset: Long,
    val type: String,
    val blobHeaderSize: Int,
    val blobSize: Int,
    val rawSize: Int?,
    val compression: PbfCompression,
)

data class PbfBlockIndex(
    val sourceSize: Long,
    val sourceLastModifiedMillis: Long,
    val sourceSha256: String,
    val entries: List<PbfBlockIndexEntry>,
) {
    val dataBlockCount: Int
        get() = entries.count { it.type == "OSMData" }
}

data class PbfScanProgress(
    val bytesRead: Long,
    val totalBytes: Long,
    val blocksRead: Int,
)

class PbfScanCancelledException : IOException("Der PBF-Blockscan wurde abgebrochen.")

object PbfBlockScanner {
    private const val MAX_BLOB_HEADER_SIZE = 64 * 1024
    private const val MAX_BLOB_SIZE = 32 * 1024 * 1024

    fun scan(
        path: Path,
        onProgress: (PbfScanProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): PbfBlockIndex {
        val normalizedPath = path.toAbsolutePath().normalize()
        val sourceSize = normalizedPath.fileSize()
        val lastModified = normalizedPath.getLastModifiedTime().toMillis()
        val digest = MessageDigest.getInstance("SHA-256")
        val entries = mutableListOf<PbfBlockIndexEntry>()
        var fileOffset = 0L

        Files.newInputStream(normalizedPath).use { fileInput ->
            DigestInputStream(fileInput, digest).use { input ->
                while (true) {
                    if (isCancelled()) throw PbfScanCancelledException()
                    val blobHeaderSize = input.readBigEndianIntOrNull() ?: break
                    if (blobHeaderSize !in 1..MAX_BLOB_HEADER_SIZE) {
                        throw IOException("Ungültige PBF-BlobHeader-Größe bei Offset $fileOffset: $blobHeaderSize Byte.")
                    }
                    val blobHeader = parseBlobHeader(input.readExactly(blobHeaderSize))
                    if (blobHeader.dataSize !in 1..MAX_BLOB_SIZE) {
                        throw IOException("Ungültige PBF-Blob-Größe bei Offset $fileOffset: ${blobHeader.dataSize} Byte.")
                    }
                    val blobInfo = parseBlob(input.readExactly(blobHeader.dataSize))
                    entries += PbfBlockIndexEntry(
                        sequence = entries.size,
                        fileOffset = fileOffset,
                        type = blobHeader.type,
                        blobHeaderSize = blobHeaderSize,
                        blobSize = blobHeader.dataSize,
                        rawSize = blobInfo.rawSize,
                        compression = blobInfo.compression,
                    )
                    fileOffset += 4L + blobHeaderSize + blobHeader.dataSize
                    onProgress(PbfScanProgress(fileOffset, sourceSize, entries.size))
                }
            }
        }

        if (fileOffset != sourceSize) {
            throw IOException("Der PBF-Scan endete bei $fileOffset Byte, die Datei ist jedoch $sourceSize Byte groß.")
        }
        if (entries.firstOrNull()?.type != "OSMHeader") {
            throw IOException("Die PBF-Datei beginnt nicht mit einem OSMHeader.")
        }
        return PbfBlockIndex(
            sourceSize = sourceSize,
            sourceLastModifiedMillis = lastModified,
            sourceSha256 = digest.digest().toHexString(),
            entries = entries,
        )
    }

    private fun parseBlobHeader(bytes: ByteArray): BlobHeader {
        val reader = ProtoReader(bytes)
        var type: String? = null
        var dataSize: Int? = null
        while (reader.hasRemaining()) {
            val tag = reader.readTag()
            when (tag.fieldNumber) {
                1 -> type = reader.readString(tag)
                3 -> dataSize = reader.readInt32(tag)
                else -> reader.skip(tag)
            }
        }
        return BlobHeader(
            type = type ?: throw IOException("Im PBF-BlobHeader fehlt der Typ."),
            dataSize = dataSize ?: throw IOException("Im PBF-BlobHeader fehlt die Datengröße."),
        )
    }

    private fun parseBlob(bytes: ByteArray): BlobInfo {
        val reader = ProtoReader(bytes)
        var rawSize: Int? = null
        val compressionFields = mutableListOf<PbfCompression>()
        while (reader.hasRemaining()) {
            val tag = reader.readTag()
            when (tag.fieldNumber) {
                1 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.RAW
                }
                2 -> rawSize = reader.readInt32(tag)
                3 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.ZLIB
                }
                4 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.LZMA
                }
                5 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.BZIP2
                }
                6 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.LZ4
                }
                7 -> {
                    reader.skipLengthDelimited(tag)
                    compressionFields += PbfCompression.ZSTD
                }
                else -> reader.skip(tag)
            }
        }
        if (compressionFields.size != 1) {
            throw IOException("Ein PBF-Blob muss genau eine Nutzdaten-Kompression enthalten.")
        }
        val compression = compressionFields.single()
        if (compression != PbfCompression.RAW && (rawSize == null || rawSize !in 0..MAX_BLOB_SIZE)) {
            throw IOException("Die entpackte PBF-Blob-Größe ist ungültig oder fehlt.")
        }
        return BlobInfo(compression, rawSize)
    }

    private data class BlobHeader(val type: String, val dataSize: Int)
    private data class BlobInfo(val compression: PbfCompression, val rawSize: Int?)
}

private fun InputStream.readBigEndianIntOrNull(): Int? {
    val first = read()
    if (first < 0) return null
    val remaining = readExactly(3)
    return (first shl 24) or
        ((remaining[0].toInt() and 0xFF) shl 16) or
        ((remaining[1].toInt() and 0xFF) shl 8) or
        (remaining[2].toInt() and 0xFF)
}

private fun ByteArray.toHexString(): String = joinToString(separator = "") { byte ->
    (byte.toInt() and 0xFF).toString(16).padStart(2, '0')
}
