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

import de.arxsilex.maptools.converter.core.GeoBounds
import de.arxsilex.maptools.converter.core.mtci.MtciReader
import java.io.BufferedOutputStream
import java.io.ByteArrayOutputStream
import java.io.DataOutputStream
import java.io.IOException
import java.nio.file.AtomicMoveNotSupportedException
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.StandardCopyOption
import java.util.zip.Deflater
import java.util.zip.DeflaterOutputStream

enum class MtciPbfWritePhase {
    NODES,
    WAYS,
    RELATIONS,
}

data class MtciPbfWriteProgress(
    val phase: MtciPbfWritePhase,
    val writtenObjects: Long,
    val writtenBlocks: Int,
    val expectedObjects: Long? = null,
)

data class MtciPbfWriteReport(
    val nodeCount: Long,
    val wayCount: Long,
    val relationCount: Long,
    val dataBlockCount: Int,
    val outputBytes: Long,
    val elapsedMillis: Long,
)

object MtciPbfWriter {
    private const val NODE_BATCH_SIZE = 8_000
    private const val WAY_BATCH_SIZE = 2_000
    private const val RELATION_BATCH_SIZE = 1_000
    private const val MAX_RAW_BLOCK_SIZE = 32 * 1024 * 1024

    fun write(
        sourceDirectory: Path,
        outputPath: Path,
        bounds: GeoBounds,
        onProgress: (MtciPbfWriteProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): MtciPbfWriteReport {
        val source = sourceDirectory.toAbsolutePath().normalize()
        val output = outputPath.toAbsolutePath().normalize()
        require(Files.isDirectory(source)) { "MTCI-Quelle nicht gefunden: $source" }
        require(!Files.exists(output)) { "PBF-Ausgabe existiert bereits: $output" }
        Files.createDirectories(output.parent ?: throw IOException("PBF-Ausgabe benötigt ein Verzeichnis."))
        val temporary = output.resolveSibling("${output.fileName}.partial")
        Files.deleteIfExists(temporary)
        val startedAt = System.nanoTime()
        var nodeCount = 0L
        var wayCount = 0L
        var relationCount = 0L
        var dataBlocks = 0

        try {
            DataOutputStream(BufferedOutputStream(Files.newOutputStream(temporary), 256 * 1024)).use { stream ->
                stream.writeFileBlock("OSMHeader", encodeHeader(bounds))

                val nodes = ArrayList<OsmPreviewNode>(NODE_BATCH_SIZE)
                fun flushNodes() {
                    if (nodes.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeNodes(nodes))
                    nodeCount += nodes.size
                    dataBlocks++
                    nodes.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.NODES, nodeCount, dataBlocks))
                }
                MtciReader.forEachNode(source) { node ->
                    nodes += node
                    if (nodes.size == NODE_BATCH_SIZE) flushNodes()
                }
                flushNodes()

                val ways = ArrayList<OsmPreviewWay>(WAY_BATCH_SIZE)
                fun flushWays() {
                    if (ways.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeWays(ways))
                    wayCount += ways.size
                    dataBlocks++
                    ways.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.WAYS, wayCount, dataBlocks))
                }
                MtciReader.forEachWay(source) { way ->
                    ways += way
                    if (ways.size == WAY_BATCH_SIZE) flushWays()
                }
                flushWays()

                val relations = ArrayList<OsmPreviewRelation>(RELATION_BATCH_SIZE)
                fun flushRelations() {
                    if (relations.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeRelations(relations))
                    relationCount += relations.size
                    dataBlocks++
                    relations.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.RELATIONS, relationCount, dataBlocks))
                }
                MtciReader.forEachRelation(source) { relation ->
                    relations += relation
                    if (relations.size == RELATION_BATCH_SIZE) flushRelations()
                }
                flushRelations()
            }
            try {
                Files.move(temporary, output, StandardCopyOption.ATOMIC_MOVE)
            } catch (_: AtomicMoveNotSupportedException) {
                Files.move(temporary, output)
            }
            return MtciPbfWriteReport(
                nodeCount = nodeCount,
                wayCount = wayCount,
                relationCount = relationCount,
                dataBlockCount = dataBlocks,
                outputBytes = Files.size(output),
                elapsedMillis = (System.nanoTime() - startedAt) / 1_000_000L,
            )
        } catch (exception: Throwable) {
            Files.deleteIfExists(temporary)
            throw exception
        }
    }

    /**
     * Writes an ordered primitive stream without requiring an MTCI directory. This is used for
     * vector sources such as Mapsforge files whose original OSM identifiers are no longer present.
     */
    fun writePrimitives(
        outputPath: Path,
        bounds: GeoBounds,
        nodes: Sequence<OsmPreviewNode>,
        ways: Sequence<OsmPreviewWay>,
        relations: Sequence<OsmPreviewRelation> = emptySequence(),
        onProgress: (MtciPbfWriteProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): MtciPbfWriteReport {
        val output = outputPath.toAbsolutePath().normalize()
        require(!Files.exists(output)) { "PBF-Ausgabe existiert bereits: $output" }
        Files.createDirectories(output.parent ?: throw IOException("PBF-Ausgabe benötigt ein Verzeichnis."))
        val temporary = output.resolveSibling("${output.fileName}.partial")
        Files.deleteIfExists(temporary)
        val startedAt = System.nanoTime()
        var nodeCount = 0L
        var wayCount = 0L
        var relationCount = 0L
        var dataBlocks = 0

        try {
            DataOutputStream(BufferedOutputStream(Files.newOutputStream(temporary), 256 * 1024)).use { stream ->
                stream.writeFileBlock("OSMHeader", encodeHeader(bounds))

                val nodeBatch = ArrayList<OsmPreviewNode>(NODE_BATCH_SIZE)
                fun flushNodes() {
                    if (nodeBatch.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeNodes(nodeBatch))
                    nodeCount += nodeBatch.size
                    dataBlocks++
                    nodeBatch.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.NODES, nodeCount, dataBlocks))
                }
                nodes.forEach { node ->
                    nodeBatch += node
                    if (nodeBatch.size == NODE_BATCH_SIZE) flushNodes()
                }
                flushNodes()

                val wayBatch = ArrayList<OsmPreviewWay>(WAY_BATCH_SIZE)
                fun flushWays() {
                    if (wayBatch.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeWays(wayBatch))
                    wayCount += wayBatch.size
                    dataBlocks++
                    wayBatch.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.WAYS, wayCount, dataBlocks))
                }
                ways.forEach { way ->
                    wayBatch += way
                    if (wayBatch.size == WAY_BATCH_SIZE) flushWays()
                }
                flushWays()

                val relationBatch = ArrayList<OsmPreviewRelation>(RELATION_BATCH_SIZE)
                fun flushRelations() {
                    if (relationBatch.isEmpty()) return
                    checkCancellation(isCancelled)
                    stream.writeFileBlock("OSMData", encodeRelations(relationBatch))
                    relationCount += relationBatch.size
                    dataBlocks++
                    relationBatch.clear()
                    onProgress(MtciPbfWriteProgress(MtciPbfWritePhase.RELATIONS, relationCount, dataBlocks))
                }
                relations.forEach { relation ->
                    relationBatch += relation
                    if (relationBatch.size == RELATION_BATCH_SIZE) flushRelations()
                }
                flushRelations()
            }
            require(nodeCount + wayCount + relationCount > 0L) {
                "Die ausgewählte Mapsforge-Quelle enthält im Ausschnitt keine exportierbaren Objekte."
            }
            try {
                Files.move(temporary, output, StandardCopyOption.ATOMIC_MOVE)
            } catch (_: AtomicMoveNotSupportedException) {
                Files.move(temporary, output)
            }
            return MtciPbfWriteReport(
                nodeCount = nodeCount,
                wayCount = wayCount,
                relationCount = relationCount,
                dataBlockCount = dataBlocks,
                outputBytes = Files.size(output),
                elapsedMillis = (System.nanoTime() - startedAt) / 1_000_000L,
            )
        } catch (exception: Throwable) {
            Files.deleteIfExists(temporary)
            throw exception
        }
    }

    private fun encodeHeader(bounds: GeoBounds): ByteArray {
        val bbox = PbfProtoWriter().apply {
            sint64(1, Math.round(bounds.west * 1_000_000_000.0))
            sint64(2, Math.round(bounds.east * 1_000_000_000.0))
            sint64(3, Math.round(bounds.north * 1_000_000_000.0))
            sint64(4, Math.round(bounds.south * 1_000_000_000.0))
        }.toByteArray()
        return PbfProtoWriter().apply {
            bytes(1, bbox)
            string(4, "OsmSchema-V0.6")
            string(4, "DenseNodes")
            string(16, "ArxSilex MapTools Converter")
            string(17, "OpenStreetMap")
        }.toByteArray()
    }

    private fun encodeNodes(nodes: List<OsmPreviewNode>): ByteArray {
        val strings = StringTableBuilder().apply { nodes.forEach { addTags(it.tags) } }
        var id = 0L
        var latitude = 0L
        var longitude = 0L
        val idDeltas = ArrayList<Long>(nodes.size)
        val latitudeDeltas = ArrayList<Long>(nodes.size)
        val longitudeDeltas = ArrayList<Long>(nodes.size)
        val keysValues = ArrayList<Int>()
        nodes.forEach { node ->
            val encodedLatitude = coordinate(node.latitude)
            val encodedLongitude = coordinate(node.longitude)
            idDeltas += node.id - id
            latitudeDeltas += encodedLatitude - latitude
            longitudeDeltas += encodedLongitude - longitude
            id = node.id
            latitude = encodedLatitude
            longitude = encodedLongitude
            node.tags.entries.sortedWith(compareBy({ it.key }, { it.value })).forEach { (key, value) ->
                keysValues += strings.indexOf(key)
                keysValues += strings.indexOf(value)
            }
            keysValues += 0
        }
        val dense = PbfProtoWriter().apply {
            packedSInt64(1, idDeltas)
            packedSInt64(8, latitudeDeltas)
            packedSInt64(9, longitudeDeltas)
            packedInt32(10, keysValues)
        }.toByteArray()
        val group = PbfProtoWriter().apply { bytes(2, dense) }.toByteArray()
        return primitiveBlock(strings, group)
    }

    private fun encodeWays(ways: List<OsmPreviewWay>): ByteArray {
        val strings = StringTableBuilder().apply { ways.forEach { addTags(it.tags) } }
        val group = PbfProtoWriter().apply {
            ways.forEach { way ->
                val keys = way.tags.entries.sortedWith(compareBy({ it.key }, { it.value }))
                var reference = 0L
                val referenceDeltas = way.nodeIds.map { nodeId ->
                    (nodeId - reference).also { reference = nodeId }
                }
                bytes(
                    3,
                    PbfProtoWriter().apply {
                        int64(1, way.id)
                        packedUInt32(2, keys.map { strings.indexOf(it.key) })
                        packedUInt32(3, keys.map { strings.indexOf(it.value) })
                        packedSInt64(8, referenceDeltas)
                    }.toByteArray(),
                )
            }
        }.toByteArray()
        return primitiveBlock(strings, group)
    }

    private fun encodeRelations(relations: List<OsmPreviewRelation>): ByteArray {
        val strings = StringTableBuilder().apply {
            relations.forEach { relation ->
                addTags(relation.tags)
                relation.members.forEach { add(it.role) }
            }
        }
        val group = PbfProtoWriter().apply {
            relations.forEach { relation ->
                val keys = relation.tags.entries.sortedWith(compareBy({ it.key }, { it.value }))
                var memberId = 0L
                val memberDeltas = relation.members.map { member ->
                    (member.id - memberId).also { memberId = member.id }
                }
                bytes(
                    4,
                    PbfProtoWriter().apply {
                        int64(1, relation.id)
                        packedUInt32(2, keys.map { strings.indexOf(it.key) })
                        packedUInt32(3, keys.map { strings.indexOf(it.value) })
                        packedInt32(8, relation.members.map { strings.indexOf(it.role) })
                        packedSInt64(9, memberDeltas)
                        packedInt32(10, relation.members.map { it.type.ordinal })
                    }.toByteArray(),
                )
            }
        }.toByteArray()
        return primitiveBlock(strings, group)
    }

    private fun primitiveBlock(strings: StringTableBuilder, group: ByteArray): ByteArray {
        val stringTable = PbfProtoWriter().apply {
            strings.values.forEach { string(1, it) }
        }.toByteArray()
        return PbfProtoWriter().apply {
            bytes(1, stringTable)
            bytes(2, group)
            int32(17, 100)
        }.toByteArray().also { bytes ->
            if (bytes.size > MAX_RAW_BLOCK_SIZE) throw IOException("MTCI-PBF-Datenblock überschreitet 32 MiB.")
        }
    }

    private fun coordinate(value: Double): Long {
        // PrimitiveBlock uses granularity 100, i.e. one stored unit equals 10^-7 degrees. Reader
        // formats such as Mapsforge can produce interpolated doubles between those units, so the
        // nearest representable coordinate is used instead of rejecting an otherwise valid shape.
        return Math.round(value * 10_000_000.0)
    }

    private fun DataOutputStream.writeFileBlock(type: String, rawData: ByteArray) {
        val compressed = rawData.zlib()
        val blob = PbfProtoWriter().apply {
            int32(2, rawData.size)
            bytes(3, compressed)
        }.toByteArray()
        val header = PbfProtoWriter().apply {
            string(1, type)
            int32(3, blob.size)
        }.toByteArray()
        if (header.size > 64 * 1024 || blob.size > MAX_RAW_BLOCK_SIZE) {
            throw IOException("Erzeugter PBF-FileBlock überschreitet die Formatgrenze.")
        }
        writeInt(header.size)
        write(header)
        write(blob)
    }

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

    private fun checkCancellation(isCancelled: () -> Boolean) {
        if (isCancelled()) throw PbfScanCancelledException()
    }
}

private class StringTableBuilder {
    private val indexes = linkedMapOf("" to 0)
    val values: List<String>
        get() = indexes.keys.toList()

    fun add(value: String) {
        indexes.getOrPut(value) { indexes.size }
    }

    fun addTags(tags: Map<String, String>) {
        tags.entries.sortedWith(compareBy({ it.key }, { it.value })).forEach { (key, value) ->
            add(key)
            add(value)
        }
    }

    fun indexOf(value: String): Int = indexes[value] ?: error("String fehlt in der PBF-StringTable: $value")
}
