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

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.UUID

data class MtciMaterializationProgress(
    val phase: MtciSpatialPhase,
    val writtenObjects: Long,
    val expectedObjects: Long,
)

data class MtciMaterializationReport(
    val nodeCount: Long,
    val wayCount: Long,
    val relationCount: Long,
    val missingSourceNodes: Int,
    val missingSourceWays: Int,
    val missingSourceRelations: Int,
    val omittedRelationMembers: Long,
    val outputBytes: Long,
    val elapsedMillis: Long,
    val files: List<MtciFileDescriptor>,
)

object MtciSelectionMaterializer {
    fun materialize(
        sourceDirectory: Path,
        selection: MtciSpatialSelection,
        outputDirectory: Path,
        shardCount: Int = 64,
        onProgress: (MtciMaterializationProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): MtciMaterializationReport {
        require(shardCount in 1..1024 && shardCount.countOneBits() == 1) {
            "Die MTCI-Shardanzahl muss eine Zweierpotenz zwischen 1 und 1024 sein."
        }
        val source = sourceDirectory.toAbsolutePath().normalize()
        val output = outputDirectory.toAbsolutePath().normalize()
        require(Files.isDirectory(source)) { "MTCI-Quelle 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 startedAt = System.nanoTime()
        var nodesWritten = 0L
        var waysWritten = 0L
        var relationsWritten = 0L
        var omittedRelationMembers = 0L
        val missingNodes = HashSet(selection.requiredNodeIds)
        val missingWays = HashSet(selection.selectedWayIds)
        val missingRelations = HashSet(selection.selectedRelationIds)

        try {
            val descriptors = MtciShardWriter(staging, shardCount).use { writer ->
                MtciReader.forEachNode(source) { node ->
                    if (node.id in selection.requiredNodeIds) {
                        writer.writeNode(node)
                        nodesWritten++
                        missingNodes.remove(node.id)
                        if (nodesWritten % 100_000L == 0L) {
                            checkCancellation(isCancelled)
                            onProgress(
                                MtciMaterializationProgress(
                                    MtciSpatialPhase.NODES,
                                    nodesWritten,
                                    selection.requiredNodeIds.size.toLong(),
                                ),
                            )
                        }
                    }
                }
                MtciReader.forEachWay(source) { way ->
                    if (way.id in selection.selectedWayIds) {
                        val missingWayNode = way.nodeIds.firstOrNull(missingNodes::contains)
                        if (missingWayNode != null) {
                            throw IOException(
                                "Way ${way.id} verweist auf Node $missingWayNode, die in der MTCI-Quelle fehlt.",
                            )
                        }
                        writer.writeWay(way)
                        waysWritten++
                        missingWays.remove(way.id)
                        if (waysWritten % 50_000L == 0L) {
                            checkCancellation(isCancelled)
                            onProgress(
                                MtciMaterializationProgress(
                                    MtciSpatialPhase.WAYS,
                                    waysWritten,
                                    selection.selectedWayIds.size.toLong(),
                                ),
                            )
                        }
                    }
                }
                MtciReader.forEachRelation(source) { relation ->
                    if (relation.id in selection.selectedRelationIds) {
                        val availableMembers = relation.members.filter { member ->
                            when (member.type) {
                                de.arxsilex.maptools.converter.core.pbf.OsmRelationMemberType.NODE ->
                                    member.id in selection.requiredNodeIds && member.id !in missingNodes
                                de.arxsilex.maptools.converter.core.pbf.OsmRelationMemberType.WAY ->
                                    member.id in selection.selectedWayIds
                                de.arxsilex.maptools.converter.core.pbf.OsmRelationMemberType.RELATION ->
                                    member.id in selection.selectedRelationIds
                            }
                        }
                        omittedRelationMembers += relation.members.size - availableMembers.size
                        writer.writeRelation(relation.copy(members = availableMembers))
                        relationsWritten++
                        missingRelations.remove(relation.id)
                    }
                }
                checkCancellation(isCancelled)
                writer.finish()
            }
            val elapsedMillis = (System.nanoTime() - startedAt) / 1_000_000L
            MtciIntegrityValidator.validate(staging).requireValid()
            val parentManifest = source.resolve("manifest.json")
            val selectionManifest = selectionManifestJson(
                selection = selection,
                parentManifestSha256 = parentManifest.takeIf(Files::isRegularFile)?.sha256() ?: "",
                descriptors = descriptors,
                nodesWritten = nodesWritten,
                waysWritten = waysWritten,
                relationsWritten = relationsWritten,
                missingSourceNodes = missingNodes.size,
                missingSourceWays = missingWays.size,
                missingSourceRelations = missingRelations.size,
                omittedRelationMembers = omittedRelationMembers,
            )
            Files.writeString(staging.resolve("manifest.json"), selectionManifest)
            try {
                Files.move(staging, output, StandardCopyOption.ATOMIC_MOVE)
            } catch (_: AtomicMoveNotSupportedException) {
                Files.move(staging, output)
            }
            val outputBytes = descriptors.sumOf(MtciFileDescriptor::bytes) + Files.size(output.resolve("manifest.json"))
            return MtciMaterializationReport(
                nodeCount = nodesWritten,
                wayCount = waysWritten,
                relationCount = relationsWritten,
                missingSourceNodes = missingNodes.size,
                missingSourceWays = missingWays.size,
                missingSourceRelations = missingRelations.size,
                omittedRelationMembers = omittedRelationMembers,
                outputBytes = outputBytes,
                elapsedMillis = elapsedMillis,
                files = descriptors,
            )
        } catch (exception: Throwable) {
            deleteMtciTree(staging)
            throw exception
        }
    }

    private fun selectionManifestJson(
        selection: MtciSpatialSelection,
        parentManifestSha256: String,
        descriptors: List<MtciFileDescriptor>,
        nodesWritten: Long,
        waysWritten: Long,
        relationsWritten: Long,
        missingSourceNodes: Int,
        missingSourceWays: Int,
        missingSourceRelations: Int,
        omittedRelationMembers: Long,
    ): String = buildString {
        appendLine("{")
        appendLine("  \"schemaVersion\": $MTCI_SCHEMA_VERSION,")
        appendLine("  \"kind\": \"spatial-selection\",")
        appendLine("  \"byteOrder\": \"BIG_ENDIAN\",")
        appendLine("  \"parentManifestSha256\": \"$parentManifestSha256\",")
        appendLine(
            "  \"bounds\": {\"west\": ${selection.bounds.west}, \"south\": ${selection.bounds.south}, " +
                "\"east\": ${selection.bounds.east}, \"north\": ${selection.bounds.north}},",
        )
        appendLine("  \"nodeCount\": $nodesWritten,")
        appendLine("  \"wayCount\": $waysWritten,")
        appendLine("  \"relationCount\": $relationsWritten,")
        appendLine("  \"missingSourceNodes\": $missingSourceNodes,")
        appendLine("  \"missingSourceWays\": $missingSourceWays,")
        appendLine("  \"missingSourceRelations\": $missingSourceRelations,")
        appendLine("  \"omittedRelationMembers\": $omittedRelationMembers,")
        appendLine("  \"files\": [")
        descriptors.forEachIndexed { index, file ->
            append("    {\"name\": \"${file.name}\", \"type\": \"${file.type}\", ")
            append("\"shard\": ${file.shard}, \"records\": ${file.records}, ")
            append("\"bytes\": ${file.bytes}, \"sha256\": \"${file.sha256}\"}")
            if (index != descriptors.lastIndex) append(',')
            appendLine()
        }
        appendLine("  ]")
        appendLine("}")
    }

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

private fun Path.sha256(): String {
    val digest = MessageDigest.getInstance("SHA-256")
    Files.newInputStream(this).use { input ->
        val buffer = ByteArray(64 * 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')
    }
}
