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

import de.arxsilex.maptools.converter.core.GeoBounds
import de.arxsilex.maptools.converter.core.GeoCircleRegion
import de.arxsilex.maptools.converter.core.pbf.ParallelPbfDecodeCancelledException
import de.arxsilex.maptools.converter.core.pbf.ParallelPbfDecodeOptions
import de.arxsilex.maptools.converter.core.pbf.PbfBlockScanner
import de.arxsilex.maptools.converter.core.pbf.PbfDataBlockReader
import de.arxsilex.maptools.converter.core.pbf.PbfHeaderReader
import de.arxsilex.maptools.converter.core.pbf.PrimitiveBlockDecoder
import de.arxsilex.maptools.converter.core.pbf.MtciPbfWriter
import de.arxsilex.maptools.converter.core.pbf.decoderPbfBlock
import de.arxsilex.maptools.converter.core.pbf.samplePrimitiveBlock
import java.io.DataInputStream
import java.nio.file.Files
import java.nio.file.Path
import java.util.Comparator
import kotlin.io.path.createTempDirectory
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertFalse
import kotlin.test.assertTrue

class MtciIntermediateBuilderTest {
    @Test
    fun `writes deterministic checksummed shards independent of worker count`() {
        val parent = createTempDirectory("mtci-test-")
        val source = parent.resolve("fixture.osm.pbf")
        Files.write(
            source,
            decoderPbfBlock("OSMHeader", byteArrayOf(1), compressed = false) +
                decoderPbfBlock("OSMData", samplePrimitiveBlock(), compressed = true) +
                decoderPbfBlock("OSMData", samplePrimitiveBlock(), compressed = true),
        )
        val index = PbfBlockScanner.scan(source)
        val sequential = parent.resolve("sequential")
        val parallel = parent.resolve("parallel")

        try {
            val first = MtciIntermediateBuilder.build(
                source,
                index,
                sequential,
                options = MtciBuildOptions(
                    shardCount = 4,
                    parallel = ParallelPbfDecodeOptions(workerCount = 1, maximumBufferedBlocks = 2),
                ),
            )
            val second = MtciIntermediateBuilder.build(
                source,
                index,
                parallel,
                options = MtciBuildOptions(
                    shardCount = 4,
                    parallel = ParallelPbfDecodeOptions(workerCount = 2, maximumBufferedBlocks = 4),
                ),
            )

            assertEquals(4, first.nodeCount)
            assertEquals(2, first.wayCount)
            assertEquals(2, first.relationCount)
            assertEquals(first, second)
            assertEquals(relativeFiles(sequential), relativeFiles(parallel))
            relativeFiles(sequential).forEach { relative ->
                assertContentEquals(Files.readAllBytes(sequential.resolve(relative)), Files.readAllBytes(parallel.resolve(relative)))
            }
            assertTrue(first.files.all { it.sha256.length == 64 })
            assertEquals(12, first.files.size)
            assertShardHeader(sequential.resolve("nodes-000.mtcn"), MTCI_NODE_MAGIC, 0, 4)
            assertTrue(Files.readString(sequential.resolve("manifest.json")).contains("\"nodeCount\": 4"))
            val nodes = mutableListOf<Long>()
            val ways = mutableListOf<Long>()
            val relations = mutableListOf<Long>()
            val readStatistics = MtciReader.readAll(
                sequential,
                onNode = { nodes += it.id },
                onWay = { ways += it.id },
                onRelation = { relations += it.id },
            )
            assertEquals(MtciReadStatistics(4, 2, 2), readStatistics)
            assertEquals(listOf(100L, 100L, 105L, 105L), nodes.sorted())
            assertEquals(listOf(200L, 200L), ways.sorted())
            assertEquals(listOf(300L, 300L), relations.sorted())
            val selection = MtciSpatialSelector.analyze(
                sequential,
                GeoBounds(west = 7.99999, south = 49.99999, east = 8.00001, north = 50.00001),
            )
            assertEquals(setOf(100L), selection.nodesInsideBounds)
            assertEquals(setOf(100L, 105L), selection.requiredNodeIds)
            assertEquals(setOf(200L), selection.selectedWayIds)
            assertEquals(setOf(300L), selection.selectedRelationIds)
            val disjointSelection = MtciSpatialSelector.analyze(
                sequential,
                listOf(
                    GeoBounds(7.99999, 49.99999, 8.00001, 50.00001),
                    GeoBounds(8.00019, 50.00009, 8.00021, 50.00011),
                ),
            )
            assertEquals(setOf(100L, 105L), disjointSelection.nodesInsideBounds)
            assertEquals(GeoBounds(7.99999, 49.99999, 8.00021, 50.00011), disjointSelection.bounds)
            val circleSelection = MtciSpatialSelector.analyze(
                sequential,
                listOf(GeoCircleRegion(50.0, 8.0, radiusKilometers = 0.005)),
            )
            assertEquals(setOf(100L), circleSelection.nodesInsideBounds)
        } finally {
            deleteTestTree(parent)
        }
    }

    @Test
    fun `cancellation removes staging directory and leaves no final output`() {
        val parent = createTempDirectory("mtci-cancel-")
        val source = parent.resolve("fixture.osm.pbf")
        Files.write(
            source,
            decoderPbfBlock("OSMHeader", byteArrayOf(1), compressed = false) +
                decoderPbfBlock("OSMData", samplePrimitiveBlock(), compressed = true),
        )
        val output = parent.resolve("cancelled")

        try {
            assertFailsWith<ParallelPbfDecodeCancelledException> {
                MtciIntermediateBuilder.build(
                    source,
                    PbfBlockScanner.scan(source),
                    output,
                    options = MtciBuildOptions(
                        shardCount = 2,
                        parallel = ParallelPbfDecodeOptions(workerCount = 1, maximumBufferedBlocks = 2),
                    ),
                    isCancelled = { true },
                )
            }
            assertFalse(Files.exists(output))
            Files.list(parent).use { children ->
                assertFalse(children.anyMatch { it.fileName.toString().startsWith(".cancelled.partial-") })
            }
        } finally {
            deleteTestTree(parent)
        }
    }

    @Test
    fun `materializes a referentially complete spatial selection`() {
        val parent = createTempDirectory("mtci-materialize-")
        val source = parent.resolve("fixture.osm.pbf")
        Files.write(
            source,
            decoderPbfBlock("OSMHeader", byteArrayOf(1), compressed = false) +
                decoderPbfBlock("OSMData", samplePrimitiveBlock(), compressed = true),
        )
        val full = parent.resolve("full")
        val selected = parent.resolve("selected")

        try {
            MtciIntermediateBuilder.build(
                source,
                PbfBlockScanner.scan(source),
                full,
                options = MtciBuildOptions(
                    shardCount = 2,
                    parallel = ParallelPbfDecodeOptions(workerCount = 1, maximumBufferedBlocks = 2),
                ),
            )
            val selection = MtciSpatialSelector.analyze(
                full,
                GeoBounds(west = 7.99999, south = 49.99999, east = 8.00001, north = 50.00001),
            )
            val report = MtciSelectionMaterializer.materialize(
                full,
                selection,
                selected,
                shardCount = 2,
            )

            assertEquals(2, report.nodeCount)
            assertEquals(1, report.wayCount)
            assertEquals(1, report.relationCount)
            assertEquals(0, report.missingSourceNodes)
            assertEquals(0, report.missingSourceWays)
            assertEquals(0, report.missingSourceRelations)
            assertEquals(0, report.omittedRelationMembers)
            assertEquals(MtciReadStatistics(2, 1, 1), MtciReader.readAll(selected))
            assertTrue(MtciIntegrityValidator.validate(selected).valid)
            val pbfOutput = parent.resolve("selected.osm.pbf")
            val selectedBounds = GeoBounds(west = 7.99999, south = 49.99999, east = 8.00001, north = 50.00001)
            val pbfReport = MtciPbfWriter.write(selected, pbfOutput, selectedBounds)
            assertEquals(2, pbfReport.nodeCount)
            assertEquals(1, pbfReport.wayCount)
            assertEquals(1, pbfReport.relationCount)
            assertEquals(selectedBounds, PbfHeaderReader.read(pbfOutput).bounds)
            val pbfIndex = PbfBlockScanner.scan(pbfOutput)
            var pbfNodes = 0
            var pbfWays = 0
            var pbfRelations = 0
            PbfDataBlockReader.open(pbfOutput).use { blocks ->
                pbfIndex.entries.filter { it.type == "OSMData" }.forEach { entry ->
                    val statistics = PrimitiveBlockDecoder.decode(blocks.read(entry))
                    pbfNodes += statistics.nodeCount
                    pbfWays += statistics.wayCount
                    pbfRelations += statistics.relationCount
                }
            }
            assertEquals(2, pbfNodes)
            assertEquals(1, pbfWays)
            assertEquals(1, pbfRelations)
        } finally {
            deleteTestTree(parent)
        }
    }

    private fun assertShardHeader(path: Path, magic: Int, shard: Int, shardCount: Int) {
        DataInputStream(Files.newInputStream(path)).use { input ->
            assertEquals(magic, input.readInt())
            assertEquals(MTCI_SCHEMA_VERSION, input.readInt())
            assertEquals(shard, input.readInt())
            assertEquals(shardCount, input.readInt())
        }
    }

    private fun relativeFiles(directory: Path): List<String> = Files.list(directory).use { files ->
        files.map { it.fileName.toString() }.sorted().toList()
    }

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