package com.example.studentfaceregistry

import android.graphics.Bitmap
import android.graphics.Matrix
import android.media.Image
import androidx.camera.core.ExperimentalGetImage
import androidx.camera.core.ImageProxy

@ExperimentalGetImage
object ImageUtils {
    fun imageProxyToBitmap(imageProxy: ImageProxy, applyRotation: Boolean = true): Bitmap {
        val image = imageProxy.image ?: error("ImageProxy does not contain an image.")
        val bitmap = yuv420ToBitmap(image)
        if (!applyRotation || imageProxy.imageInfo.rotationDegrees == 0) {
            return bitmap
        }
        val matrix = Matrix().apply {
            postRotate(imageProxy.imageInfo.rotationDegrees.toFloat())
        }
        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
    }

    private fun yuv420ToBitmap(image: Image): Bitmap {
        val width = image.width
        val height = image.height
        val yPlane = image.planes[0]
        val uPlane = image.planes[1]
        val vPlane = image.planes[2]

        val yBuffer = yPlane.buffer
        val uBuffer = uPlane.buffer
        val vBuffer = vPlane.buffer
        val pixels = IntArray(width * height)

        val yRowStride = yPlane.rowStride
        val yPixelStride = yPlane.pixelStride
        val uRowStride = uPlane.rowStride
        val uPixelStride = uPlane.pixelStride
        val vRowStride = vPlane.rowStride
        val vPixelStride = vPlane.pixelStride

        var pixelIndex = 0
        for (row in 0 until height) {
            val yRow = row * yRowStride
            val uvRow = (row shr 1) * uRowStride
            val vvRow = (row shr 1) * vRowStride
            for (col in 0 until width) {
                val y = (yBuffer.get(yRow + col * yPixelStride).toInt() and 0xFF) - 16
                val u = (uBuffer.get(uvRow + (col shr 1) * uPixelStride).toInt() and 0xFF) - 128
                val v = (vBuffer.get(vvRow + (col shr 1) * vPixelStride).toInt() and 0xFF) - 128

                val y1192 = 1192 * y.coerceAtLeast(0)
                var r = y1192 + 1634 * v
                var g = y1192 - 833 * v - 400 * u
                var b = y1192 + 2066 * u

                r = r.coerceIn(0, 262143)
                g = g.coerceIn(0, 262143)
                b = b.coerceIn(0, 262143)

                pixels[pixelIndex++] = -0x1000000 or
                    ((r shl 6) and 0x00FF0000) or
                    ((g shr 2) and 0x0000FF00) or
                    ((b shr 10) and 0x000000FF)
            }
        }

        return Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888).apply {
            setPixels(pixels, 0, width, 0, 0, width, height)
        }
    }
}