/*
 * Copyright 2023 Shubham Panchal
 * Licensed under the Apache License, Version 2.0 (the "License");
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.example.absenkuv3.facex

import android.graphics.Bitmap
import android.graphics.Rect
import com.example.absenkuv3.facex.model.FaceNetModel
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.face.FaceDetection
import com.google.mlkit.vision.face.FaceDetectorOptions
import kotlinx.coroutines.*
import kotlinx.coroutines.tasks.await

// Utility class to read images from internal storage
class FileReader(private var faceNetModel: FaceNetModel) {

    private val realTimeOpts = FaceDetectorOptions.Builder()
        .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_FAST)
        .build()
    private val mainScope = CoroutineScope(Dispatchers.Main)
    private var numImagesWithNoFaces = 0
    private var imageCounter = 0
    private var numImages = 0
    private var data = ArrayList<Pair<String, Bitmap>>()
    private lateinit var callback: ProcessCallback

    // imageData will be provided to the MainActivity via ProcessCallback ( see the run() method below ) and finally,
    // used by the FrameAnalyser class.
    private val imageData = ArrayList<Pair<String, FloatArray>>()

    private val detector = FaceDetection.getClient(realTimeOpts)


    // Given the Bitmaps, extract face embeddings from then and deliver the processed embedding to ProcessCallback.


    interface ProcessCallback {
        fun onProcessCompleted(data: ArrayList<Pair<String, FloatArray>>, numImagesWithNoFaces: Int)
    }


    // Crop faces and produce embeddings ( using FaceNet ) from given image.
    // Store the embedding in imageData
    fun scanImage2(name: String, image: FloatArray, callback: ProcessCallback) {
        this.callback = callback
        mainScope.launch {

            imageData.add(Pair(name, image))
            callback.onProcessCompleted(imageData, 1)
            reset()
        }
    }

    fun drawMatrix(image: Bitmap): List<Float> {
        return runBlocking {
            val inputImage = InputImage.fromByteArray(
                BitmapUtils.bitmapToNV21ByteArray(image),
                image.width,
                image.height,
                0,
                InputImage.IMAGE_FORMAT_NV21
            )

            val faces = detector.process(inputImage).await()
            if (faces.size != 0) {
                val embedding : FloatArray
                try {
                    embedding = getEmbedding(image, faces[0].boundingBox)
                    embedding.toList()
                }
                catch (e: Exception) {
                    return@runBlocking emptyList<Float>()
                }

            } else {
                emptyList()
            }
        }
    }

    private suspend fun getEmbedding(image: Bitmap, bbox: Rect): FloatArray =
        withContext(Dispatchers.Default) {
            return@withContext faceNetModel.getFaceEmbedding(
                BitmapUtils.cropRectFromBitmap(
                    image,
                    bbox
                )
            )
        }


    private fun reset() {
        imageCounter = 0
        numImages = 0
        numImagesWithNoFaces = 0
        data.clear()
    }

}