package network.rs485.ben.computervision

import io.ktor.utils.io.pool.*
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.opencv.core.*
import org.opencv.imgcodecs.Imgcodecs
import org.opencv.imgproc.Imgproc
import java.nio.file.Path

data class MatInfo(val rows: Int, val cols: Int, val type: Int) {
    companion object {
        fun from(mat: Mat) = MatInfo(mat.rows(), mat.cols(), mat.type())
        fun from(mat: Mat, type: Int) = MatInfo(mat.rows(), mat.cols(), type)
    }
    fun create(): Mat = Mat(rows, cols, type)
}

private val histMatInfo = MatInfo(256, 1, CvType.CV_32FC1)
private val predefinedChannels = listOf(16, 24, 64, 128)
private val matPoolsToCreate = listOf(
    MatInfo(32, 32, CvType.CV_8UC1) to Runtime.getRuntime().availableProcessors() * 4,
    histMatInfo to Runtime.getRuntime().availableProcessors() * 4,
) + predefinedChannels.map {
    MatInfo(it, it, CvType.CV_8UC1) to Runtime.getRuntime().availableProcessors() * 2
}

object CvContext {
    lateinit var gaussianKernel: Mat
    lateinit var histChannelPar: MatOfInt
    lateinit var histSizePar: MatOfInt
    lateinit var histRangesPar: MatOfFloat
    lateinit var matPools: HashMap<MatInfo, MatPool>

    fun initialize() {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME)
        this.gaussianKernel = Imgproc.getGaussianKernel(GAUSSIAN_KERNEL_SIZE, -1.0)
        this.histChannelPar = MatOfInt(0)
        this.histSizePar = MatOfInt(HIST_SIZE)
        this.histRangesPar = MatOfFloat(0f, 255f)
        this.matPools = matPoolsToCreate.associateTo(HashMap()) { (matInfo, capacity) ->
            matInfo to MatPool(matInfo, capacity)
        }
    }

    internal inline fun <T> withGrayCalcHist(image: Mat, mask: Mat, block: (Mat) -> T): T = withMat(histMatInfo) { hist ->
        Imgproc.calcHist(
            /* images = */ listOf(image),
            /* channels = */ histChannelPar,
            /* mask = */ mask,
            /* hist = */ hist,
            /* histSize = */ histSizePar,
            /* ranges = */ histRangesPar,
        )
        block(hist)
    }

    internal inline fun <T> withRotationMatrix(angle: Double, center: Point, block: (Mat) -> T): T =
        Imgproc.getRotationMatrix2D(
            /* center = */ center,
            /* angle = */ Math.toDegrees(angle),
            /* scale = */ 1.0,
        ).letAutoclean(block)

    class MatPool(
        private val info: MatInfo,
        capacity: Int,
    ) : DefaultPool<Mat>(capacity) {
        override fun produceInstance(): Mat = info.create()
    }

    internal inline fun <T> withMat(matInfo: MatInfo, block: (Mat) -> T): T {
        val upperMatInfo = if (matInfo.type == CvType.CV_8UC1) {
            predefinedChannels.find { matInfo.rows <= it && matInfo.cols <= it }
                ?.let { MatInfo(it, it, CvType.CV_8UC1) }
                ?: matInfo
        } else matInfo
        return if (matPools.containsKey(upperMatInfo)) {
            val pool = matPools[upperMatInfo]!!
            val mat = pool.borrow()
            autoclean(mat, { pool.recycle(mat) }, block)
        } else {
            println("Not pooled: $matInfo")
            matInfo.create().letAutoclean(block)
        }
    }

    internal fun recycle(image: Mat) = matPools[MatInfo.from(image)]?.recycle(image) ?: image.release()

    internal suspend fun readAndConvertImage(imagePath: Path): Pair<Mat, Mat> {
        val displayImage = withContext(Dispatchers.IO) {
            Imgcodecs.imread(imagePath.toString())
        } ?: throw RuntimeException("Could not load image from $imagePath")
        println("Loaded image $imagePath as $displayImage")
        val matInfo = MatInfo(displayImage.rows(), displayImage.cols(), CvType.CV_8UC1)
        val image = matPools.computeIfAbsent(matInfo) {
            MatPool(matInfo, Runtime.getRuntime().availableProcessors())
        }.borrow()
        Imgproc.cvtColor(displayImage, image, Imgproc.COLOR_BGR2GRAY)
        return displayImage to image
    }
}

internal inline fun <V, T> autoclean(obj: V, release: () -> Unit, block: (V) -> T): T =
    runCatching { block(obj) }.also { release() }.getOrThrow()

internal inline fun <T> Mat.letAutoclean(block: (Mat) -> T): T =
    autoclean(this, this::release, block)

internal inline fun <T> withTmpMat(block: (tmp: Mat) -> T): T = Mat().letAutoclean(block)