package de.arxsilex.maptools.converter.core.mtci

import de.arxsilex.maptools.converter.core.GeoBounds
import de.arxsilex.maptools.converter.core.pbf.BoundedOrderedParallelPipeline
import de.arxsilex.maptools.converter.core.pbf.OsmPreviewNode
import de.arxsilex.maptools.converter.core.pbf.OsmPreviewRelation
import de.arxsilex.maptools.converter.core.pbf.OsmPreviewWay
import de.arxsilex.maptools.converter.core.pbf.ParallelPbfDecodeOptions
import de.arxsilex.maptools.converter.core.pbf.PbfBlockIndex
import de.arxsilex.maptools.converter.core.pbf.PbfBlockIndexEntry
import de.arxsilex.maptools.converter.core.pbf.PbfDataBlockReader
import de.arxsilex.maptools.converter.core.pbf.PrimitiveBlockData
import de.arxsilex.maptools.converter.core.pbf.PrimitiveBlockDecoder
import java.io.BufferedOutputStream
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.security.MessageDigest
import java.util.Comparator
import java.util.UUID
import kotlin.io.path.fileSize

data class MtciBuildOptions(
    val shardCount: Int = 64,
    val parallel: ParallelPbfDecodeOptions = ParallelPbfDecodeOptions(),
) {
    init {
        require(shardCount in 1..1024 && shardCount.countOneBits() == 1) {
            "Die MTCI-Shardanzahl muss eine Zweierpotenz zwischen 1 und 1024 sein."
        }
    }
}

data class MtciBuildProgress(
    val completedBlocks: Int,
    val totalBlocks: Int,
    val nodeCount: Long,
    val wayCount: Long,
    val relationCount: Long,
)

object MtciIntermediateBuilder {
    fun build(
        sourcePath: Path,
        index: PbfBlockIndex,
        outputDirectory: Path,
        sourceBounds: GeoBounds? = null,
        options: MtciBuildOptions = MtciBuildOptions(),
        onProgress: (MtciBuildProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): MtciManifest {
        val source = sourcePath.toAbsolutePath().normalize()
        val output = outputDirectory.toAbsolutePath().normalize()
        require(Files.isRegularFile(source)) { "PBF-Quelldatei nicht gefunden: $source" }
        if (Files.exists(output)) {
            require(Files.isDirectory(output) && isMtciDirectoryEmpty(output)) {
                "MTCI-Ausgabe existiert bereits und ist nicht leer: $output"
            }
            Files.delete(output)
        }
        val parent = output.parent ?: throw IOException("MTCI-Ausgabe benötigt ein übergeordnetes Verzeichnis.")
        Files.createDirectories(parent)
        val staging = parent.resolve(".${output.fileName}.partial-${UUID.randomUUID()}")
        Files.createDirectory(staging)
        val dataEntries = index.entries.filter { it.type == "OSMData" }
        var nodeCount = 0L
        var wayCount = 0L
        var relationCount = 0L

        try {
            val descriptors = MtciShardWriter(staging, options.shardCount).use { writer ->
                PbfDataBlockReader.open(source).use { blockSource ->
                    BoundedOrderedParallelPipeline.process(
                        inputs = dataEntries,
                        options = options.parallel,
                        isCancelled = isCancelled,
                        transform = { entry ->
                            DecodedMtciBlock(entry, PrimitiveBlockDecoder.decodeToData(blockSource.read(entry)))
                        },
                        onOrderedResult = { ordinal, decoded ->
                            decoded.data.nodes.forEach(writer::writeNode)
                            decoded.data.ways.forEach(writer::writeWay)
                            decoded.data.relations.forEach(writer::writeRelation)
                            nodeCount += decoded.data.nodes.size
                            wayCount += decoded.data.ways.size
                            relationCount += decoded.data.relations.size
                            onProgress(
                                MtciBuildProgress(
                                    completedBlocks = ordinal + 1,
                                    totalBlocks = dataEntries.size,
                                    nodeCount = nodeCount,
                                    wayCount = wayCount,
                                    relationCount = relationCount,
                                ),
                            )
                        },
                    )
                }
                writer.finish()
            }
            val manifest = MtciManifest(
                schemaVersion = MTCI_SCHEMA_VERSION,
                sourcePath = source.toString(),
                sourceBytes = index.sourceSize,
                sourceLastModifiedMillis = index.sourceLastModifiedMillis,
                sourceSha256 = index.sourceSha256,
                sourceBounds = sourceBounds,
                shardCount = options.shardCount,
                nodeCount = nodeCount,
                wayCount = wayCount,
                relationCount = relationCount,
                files = descriptors,
            )
            Files.writeString(staging.resolve("manifest.json"), manifest.toJson())
            try {
                Files.move(staging, output, StandardCopyOption.ATOMIC_MOVE)
            } catch (_: AtomicMoveNotSupportedException) {
                Files.move(staging, output)
            }
            return manifest
        } catch (exception: Throwable) {
            deleteMtciTree(staging)
            throw exception
        }
    }
}

private data class DecodedMtciBlock(
    val entry: PbfBlockIndexEntry,
    val data: PrimitiveBlockData,
)

internal class MtciShardWriter(
    private val directory: Path,
    private val shardCount: Int,
) : AutoCloseable {
    private val nodes = ShardOutputs(directory, "nodes", "mtcn", MTCI_NODE_MAGIC, shardCount)
    private val ways = ShardOutputs(directory, "ways", "mtcw", MTCI_WAY_MAGIC, shardCount)
    private val relations = ShardOutputs(directory, "relations", "mtcr", MTCI_RELATION_MAGIC, shardCount)
    private var closed = false

    fun writeNode(node: OsmPreviewNode) {
        val output = nodes.outputFor(node.id)
        output.writeLong(node.id)
        output.writeLong(Math.round(node.latitude * 1_000_000_000.0))
        output.writeLong(Math.round(node.longitude * 1_000_000_000.0))
        output.writeTags(node.tags)
        nodes.recordWritten(node.id)
    }

    fun writeWay(way: OsmPreviewWay) {
        val output = ways.outputFor(way.id)
        output.writeLong(way.id)
        output.writeInt(way.nodeIds.size)
        way.nodeIds.forEach(output::writeLong)
        output.writeTags(way.tags)
        ways.recordWritten(way.id)
    }

    fun writeRelation(relation: OsmPreviewRelation) {
        val output = relations.outputFor(relation.id)
        output.writeLong(relation.id)
        output.writeInt(relation.members.size)
        relation.members.forEach { member ->
            output.writeByte(member.type.ordinal)
            output.writeLong(member.id)
            output.writeUtf8(member.role)
        }
        output.writeTags(relation.tags)
        relations.recordWritten(relation.id)
    }

    fun finish(): List<MtciFileDescriptor> {
        close()
        return buildList {
            addAll(nodes.descriptors("nodes"))
            addAll(ways.descriptors("ways"))
            addAll(relations.descriptors("relations"))
        }.sortedBy(MtciFileDescriptor::name)
    }

    override fun close() {
        if (!closed) {
            closed = true
            nodes.close()
            ways.close()
            relations.close()
        }
    }
}

private class ShardOutputs(
    directory: Path,
    private val prefix: String,
    extension: String,
    magic: Int,
    private val shardCount: Int,
) : AutoCloseable {
    private val paths = List(shardCount) { shard ->
        directory.resolve("$prefix-${shard.toString().padStart(3, '0')}.$extension")
    }
    private val outputs = paths.mapIndexed { shard, path ->
        DataOutputStream(BufferedOutputStream(Files.newOutputStream(path), 64 * 1024)).apply {
            writeInt(magic)
            writeInt(MTCI_SCHEMA_VERSION)
            writeInt(shard)
            writeInt(shardCount)
        }
    }
    private val recordCounts = LongArray(shardCount)
    private var closed = false

    fun outputFor(id: Long): DataOutputStream = outputs[shardFor(id)]

    fun recordWritten(id: Long) {
        recordCounts[shardFor(id)]++
    }

    fun descriptors(type: String): List<MtciFileDescriptor> {
        check(closed) { "Shard-Dateien müssen vor der Prüfsummenbildung geschlossen sein." }
        return paths.mapIndexed { shard, path ->
            MtciFileDescriptor(
                name = path.fileName.toString(),
                type = type,
                shard = shard,
                records = recordCounts[shard],
                bytes = path.fileSize(),
                sha256 = path.sha256(),
            )
        }
    }

    override fun close() {
        if (!closed) {
            closed = true
            var firstFailure: Throwable? = null
            outputs.forEach { output ->
                runCatching(output::close).onFailure { failure ->
                    if (firstFailure == null) firstFailure = failure
                }
            }
            firstFailure?.let { throw it }
        }
    }

    private fun shardFor(id: Long): Int {
        val mixed = id xor (id ushr 33) xor (id shl 11)
        return (mixed and (shardCount - 1).toLong()).toInt()
    }
}

private fun DataOutputStream.writeTags(tags: Map<String, String>) {
    val ordered = tags.entries.sortedWith(compareBy({ it.key }, { it.value }))
    writeInt(ordered.size)
    ordered.forEach { (key, value) ->
        writeUtf8(key)
        writeUtf8(value)
    }
}

private fun DataOutputStream.writeUtf8(value: String) {
    val bytes = value.toByteArray(Charsets.UTF_8)
    require(bytes.size <= 16 * 1024 * 1024) { "MTCI-Textwert überschreitet 16 MiB." }
    writeInt(bytes.size)
    write(bytes)
}

private fun Path.sha256(): String {
    val digest = MessageDigest.getInstance("SHA-256")
    Files.newInputStream(this).use { input ->
        val buffer = ByteArray(128 * 1024)
        while (true) {
            val count = input.read(buffer)
            if (count < 0) break
            digest.update(buffer, 0, count)
        }
    }
    return digest.digest().joinToString("") { byte ->
        (byte.toInt() and 0xFF).toString(16).padStart(2, '0')
    }
}

internal fun deleteMtciTree(path: Path) {
    if (!Files.exists(path)) return
    Files.walk(path).use { paths ->
        paths.sorted(Comparator.reverseOrder()).forEach(Files::deleteIfExists)
    }
}

internal fun isMtciDirectoryEmpty(path: Path): Boolean = Files.list(path).use { children ->
    !children.findAny().isPresent
}
