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

import java.io.IOException
import java.nio.file.Path
import java.util.TreeMap
import java.util.concurrent.Callable
import java.util.concurrent.CancellationException
import java.util.concurrent.ExecutorCompletionService
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger

data class ParallelPbfDecodeOptions(
    val workerCount: Int = (Runtime.getRuntime().availableProcessors() - 1).coerceIn(1, 16),
    val maximumBufferedBlocks: Int = (workerCount * 2).coerceAtLeast(2),
) {
    init {
        require(workerCount >= 1) { "Mindestens ein Dekodier-Thread ist erforderlich." }
        require(maximumBufferedBlocks >= workerCount) {
            "Der Blockpuffer darf nicht kleiner als die Threadanzahl sein."
        }
    }
}

data class ParallelPbfDecodeProgress(
    val completedBlocks: Int,
    val totalBlocks: Int,
    val packedBytes: Long,
    val unpackedBytes: Long,
)

data class ParallelPbfDecodeReport(
    val workerCount: Int,
    val blockCount: Int,
    val nodeCount: Long,
    val wayCount: Long,
    val relationCount: Long,
    val packedBytes: Long,
    val unpackedBytes: Long,
    val elapsedMillis: Long,
    val decompressionProvider: String,
    val hardwareAccelerated: Boolean,
)

class ParallelPbfDecodeCancelledException : IOException("Die parallele PBF-Dekodierung wurde abgebrochen.")

object ParallelPbfPrimitiveDecoder {
    fun decode(
        sourcePath: Path,
        index: PbfBlockIndex,
        options: ParallelPbfDecodeOptions = ParallelPbfDecodeOptions(),
        onProgress: (ParallelPbfDecodeProgress) -> Unit = {},
        isCancelled: () -> Boolean = { false },
    ): ParallelPbfDecodeReport {
        val dataEntries = index.entries.filter { it.type == "OSMData" }
        val startedAt = System.nanoTime()
        var nodeCount = 0L
        var wayCount = 0L
        var relationCount = 0L
        var packedBytes = 0L
        var unpackedBytes = 0L
        val acceleration = PbfDecompression.accelerationStatus()

        PbfDataBlockReader.open(sourcePath).use { source ->
            BoundedOrderedParallelPipeline.process(
                inputs = dataEntries,
                options = options,
                isCancelled = isCancelled,
                transform = { entry ->
                    val bytes = source.read(entry)
                    DecodedPrimitiveBlock(
                        entry = entry,
                        unpackedBytes = bytes.size,
                        statistics = PrimitiveBlockDecoder.decode(bytes),
                    )
                },
                onOrderedResult = { ordinal, result ->
                    nodeCount += result.statistics.nodeCount
                    wayCount += result.statistics.wayCount
                    relationCount += result.statistics.relationCount
                    packedBytes += result.entry.blobSize
                    unpackedBytes += result.unpackedBytes
                    onProgress(
                        ParallelPbfDecodeProgress(
                            completedBlocks = ordinal + 1,
                            totalBlocks = dataEntries.size,
                            packedBytes = packedBytes,
                            unpackedBytes = unpackedBytes,
                        ),
                    )
                },
            )
        }

        return ParallelPbfDecodeReport(
            workerCount = options.workerCount,
            blockCount = dataEntries.size,
            nodeCount = nodeCount,
            wayCount = wayCount,
            relationCount = relationCount,
            packedBytes = packedBytes,
            unpackedBytes = unpackedBytes,
            elapsedMillis = (System.nanoTime() - startedAt) / 1_000_000L,
            decompressionProvider = acceleration.activeProvider,
            hardwareAccelerated = acceleration.accelerated,
        )
    }
}

private data class DecodedPrimitiveBlock(
    val entry: PbfBlockIndexEntry,
    val unpackedBytes: Int,
    val statistics: PrimitiveBlockStatistics,
)

internal object BoundedOrderedParallelPipeline {
    private data class OrderedResult<T>(val ordinal: Int, val value: T)

    fun <I, O> process(
        inputs: List<I>,
        options: ParallelPbfDecodeOptions,
        isCancelled: () -> Boolean,
        transform: (I) -> O,
        onOrderedResult: (Int, O) -> Unit,
    ) {
        if (inputs.isEmpty()) return
        val threadCounter = AtomicInteger()
        val executor = Executors.newFixedThreadPool(options.workerCount) { runnable ->
            Thread(runnable, "maptools-pbf-decode-${threadCounter.incrementAndGet()}").apply {
                isDaemon = true
            }
        }
        val completion = ExecutorCompletionService<OrderedResult<O>>(executor)
        val pending = TreeMap<Int, O>()
        var nextToSubmit = 0
        var nextToEmit = 0
        var running = 0

        fun fillWindow() {
            while (
                nextToSubmit < inputs.size &&
                running + pending.size < options.maximumBufferedBlocks
            ) {
                val ordinal = nextToSubmit++
                completion.submit(Callable { OrderedResult(ordinal, transform(inputs[ordinal])) })
                running++
            }
        }

        try {
            fillWindow()
            while (nextToEmit < inputs.size) {
                if (isCancelled()) throw ParallelPbfDecodeCancelledException()
                val future = completion.poll(100, TimeUnit.MILLISECONDS) ?: continue
                running--
                val completed = try {
                    future.get()
                } catch (exception: java.util.concurrent.ExecutionException) {
                    val cause = exception.cause ?: exception
                    when (cause) {
                        is IOException -> throw cause
                        is RuntimeException -> throw cause
                        is Error -> throw cause
                        else -> throw IOException("PBF-Block konnte nicht dekodiert werden.", cause)
                    }
                } catch (exception: CancellationException) {
                    throw ParallelPbfDecodeCancelledException()
                }
                pending[completed.ordinal] = completed.value
                while (true) {
                    val value = pending.remove(nextToEmit) ?: break
                    onOrderedResult(nextToEmit, value)
                    nextToEmit++
                }
                fillWindow()
            }
        } catch (exception: InterruptedException) {
            Thread.currentThread().interrupt()
            throw ParallelPbfDecodeCancelledException()
        } finally {
            executor.shutdownNow()
        }
    }
}
