package com.example.studentfaceregistry

import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.ImageFormat
import android.graphics.Matrix
import android.graphics.Rect
import android.graphics.YuvImage
import android.media.Image
import androidx.camera.core.ImageProxy
import java.io.ByteArrayOutputStream
import kotlin.math.min

object ImageUtils {
    fun imageProxyToBitmap(imageProxy: ImageProxy): Bitmap {
        val image = imageProxy.image ?: error("ImageProxy does not contain an image.")
        val nv21 = yuv420ToNv21(image)
        val yuvImage = YuvImage(nv21, ImageFormat.NV21, image.width, image.height, null)
        val output = ByteArrayOutputStream()
        yuvImage.compressToJpeg(Rect(0, 0, image.width, image.height), 90, output)
        val bitmap = BitmapFactory.decodeByteArray(output.toByteArray(), 0, output.size())
        val matrix = Matrix().apply {
            postRotate(imageProxy.imageInfo.rotationDegrees.toFloat())
        }
        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
    }

    private fun yuv420ToNv21(image: Image): ByteArray {
        val width = image.width
        val height = image.height
        val ySize = width * height
        val chromaSize = width * height / 4
        val nv21 = ByteArray(ySize + chromaSize * 2)
        copyPlane(image.planes[0], width, height, nv21, 0, 1)
        copyPlane(image.planes[2], width / 2, height / 2, nv21, ySize, 2)
        copyPlane(image.planes[1], width / 2, height / 2, nv21, ySize + 1, 2)
        return nv21
    }

    private fun copyPlane(
        plane: Image.Plane,
        width: Int,
        height: Int,
        output: ByteArray,
        offset: Int,
        outputPixelStride: Int
    ) {
        val buffer = plane.buffer
        val row = ByteArray(plane.rowStride)
        var outputIndex = offset
        for (rowIndex in 0 until height) {
            val bytesToRead = min(plane.rowStride, buffer.remaining())
            buffer.get(row, 0, bytesToRead)
            for (colIndex in 0 until width) {
                val inputIndex = colIndex * plane.pixelStride
                if (inputIndex < bytesToRead && outputIndex < output.size) {
                    output[outputIndex] = row[inputIndex]
                }
                outputIndex += outputPixelStride
            }
        }
    }
}
