package com.performancecoachlab.posedetection.camera

import android.graphics.Bitmap
import androidx.annotation.OptIn
import androidx.camera.core.ExperimentalGetImage
import androidx.camera.core.ImageProxy
import androidx.compose.ui.geometry.Rect
import com.google.android.gms.tasks.Tasks
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.pose.PoseDetector
import com.performancecoachlab.posedetection.recording.AnalysisObject
import com.performancecoachlab.posedetection.recording.AnalysisResult
import com.performancecoachlab.posedetection.recording.FrameSize
import com.performancecoachlab.posedetection.skeleton.Skeleton
import org.tensorflow.lite.support.image.TensorImage

actual enum class PlatformType {
    ANDROID, IOS;

    actual companion object {
        actual fun getCurrentPlatform(): PlatformType = ANDROID
    }
}

@OptIn(ExperimentalGetImage::class)
fun ImageProxy.process(
    objectDetector: org.tensorflow.lite.task.vision.detector.ObjectDetector?,
    poseDetector: PoseDetector?,
    timestamp: Long,
    focusArea: Rect?,
    onComplete: (AnalysisResult, Bitmap) -> Unit
) {
    val bitmap = toBitmap().rotate(imageInfo.rotationDegrees.toFloat())
    val tensorImage = TensorImage.fromBitmap(bitmap)
    val mlKitImage = InputImage.fromBitmap(bitmap.applyFocusAreaMask(focusArea,imageInfo.rotationDegrees
    ), imageInfo.rotationDegrees)
    process(
        tensorImage = tensorImage,
        mlKitImage = mlKitImage,
        objectDetector = objectDetector,
        poseDetector = poseDetector,
        timestamp = timestamp,
        width = width,
        height = height,
        bitmap = bitmap,
        onComplete = onComplete
    )
}

private fun Rect?.toGraphicsRect(width: Int, height: Int):android.graphics.Rect {
    return this?.let {
        android.graphics.Rect((it.left*width).toInt(),
            (it.top*height).toInt(),
            (it.right*width).toInt(),
            (it.bottom*height).toInt()
        )
    }?: android.graphics.Rect(0, 0, width, height)
}

fun Bitmap.process(
    objectDetector: org.tensorflow.lite.task.vision.detector.ObjectDetector?,
    poseDetector: PoseDetector,
    timestamp: Long,
    focusArea: Rect?,
    onComplete: (AnalysisResult, Bitmap) -> Unit
) {
    val tensorImage = TensorImage.fromBitmap(this)
    val mlKitImage = InputImage.fromBitmap(this.applyFocusAreaMask(focusArea), 0)
    process(
        tensorImage = tensorImage,
        mlKitImage = mlKitImage,
        objectDetector = objectDetector,
        poseDetector = poseDetector,
        timestamp = timestamp,
        width = width,
        height = height,
        bitmap = this,
        onComplete = onComplete
    )
}

/**
 * Crops the bitmap to the specified focus area rectangle
 * @param focusArea The rectangle area to crop to (in normalized coordinates 0.0-1.0)
 * @return A new bitmap cropped to the focus area, or the original bitmap if focusArea is null
 */
fun Bitmap.cropToFocusArea(focusArea: Rect?): Bitmap {
    return focusArea?.let { rect ->
        val left = (rect.left * width.toFloat()).toInt().coerceIn(0, width)
        val top = (rect.top * height.toFloat()).toInt().coerceIn(0, height)
        val right = (rect.right * width.toFloat()).toInt().coerceIn(left, width)
        val bottom = (rect.bottom * height.toFloat()).toInt().coerceIn(top, height)

        val cropWidth = right - left
        val cropHeight = bottom - top

        if (cropWidth > 0 && cropHeight > 0) {
            Bitmap.createBitmap(this, left, top, cropWidth, cropHeight)
        } else {
            this
        }
    } ?: this
}

/**
 * Creates a copy of the bitmap with everything outside the focus area blacked out
 * @param focusArea The rectangle area to keep visible (in normalized coordinates 0.0-1.0)
 * @param angle The rotation angle in degrees (must be a multiple of 90) to apply to the focus area rectangle
 * @return A new bitmap with areas outside the focus area blacked out, or the original bitmap if focusArea is null
 */
fun Bitmap.applyFocusAreaMask(focusArea: Rect?, _angle: Int = 0): Bitmap {
    return focusArea?.let { rect ->
        val result = this.copy(this.config ?: Bitmap.Config.ARGB_8888, true)
        val canvas = android.graphics.Canvas(result)
        val paint = android.graphics.Paint().apply {
            color = android.graphics.Color.BLACK
        }

        val angle = 0
        // Transform the rectangle coordinates based on the angle
        val transformedRect = when (angle % 360) {
            90 -> Rect(
                left = rect.top,
                top = 1f - rect.right,
                right = rect.bottom,
                bottom = 1f - rect.left
            )
            180 -> Rect(
                left = 1f - rect.right,
                top = 1f - rect.bottom,
                right = 1f - rect.left,
                bottom = 1f - rect.top
            )
            270 -> Rect(
                left = 1f - rect.bottom,
                top = rect.left,
                right = 1f - rect.top,
                bottom = rect.right
            )
            else -> rect // 0 degrees or any other angle
        }

        val focusRect = transformedRect.toGraphicsRect(width, height)

        // Black out top area
        if (focusRect.top > 0) {
            canvas.drawRect(0f, 0f, width.toFloat(), focusRect.top.toFloat(), paint)
        }

        // Black out bottom area
        if (focusRect.bottom < height) {
            canvas.drawRect(0f, focusRect.bottom.toFloat(), width.toFloat(),height.toFloat(), paint)
        }

        // Black out left area
        if (focusRect.left > 0) {
            canvas.drawRect(0f, focusRect.top.toFloat(), focusRect.left.toFloat(), focusRect.bottom.toFloat(), paint)
        }

        // Black out right area
        if (focusRect.right < width) {
            canvas.drawRect(focusRect.right.toFloat(), focusRect.top.toFloat(), width.toFloat(), focusRect.bottom.toFloat(), paint)
        }

        result
    } ?: this
}

private fun process(
    tensorImage: TensorImage,
    mlKitImage: InputImage?,
    objectDetector: org.tensorflow.lite.task.vision.detector.ObjectDetector?,
    poseDetector: PoseDetector?,
    timestamp: Long,
    width: Int,
    height: Int,
    bitmap: Bitmap,
    onComplete: (AnalysisResult, Bitmap) -> Unit
) {
    val objectsDetected = objectDetector?.detect(tensorImage)?.map { result ->
        AnalysisObject(
            boundingBox = result.boundingBox.let {
                Rect(
                    left = it.left,
                    top = it.top,
                    right = it.right,
                    bottom = it.bottom
                )
            },
            trackingId = 0,
            labels = result.categories.map { category ->
                com.performancecoachlab.posedetection.recording.Label(
                    category.label,
                    category.score
                )
            },
            frameSize = FrameSize(
                width = width,
                height = height
            )
        )
    } ?: emptyList()
    var skeleton: Skeleton? = null
    val poseDetectionTask = mlKitImage?.let {
        val rotation = it.rotationDegrees
        poseDetector?.process(it)?.addOnSuccessListener { pose ->
            skeleton = skeleton(pose, timestamp, width, height).let{
                it.rotate(rotation)
            }
        }?.addOnFailureListener { e ->
            //println(e)
        }
    }
    Tasks.whenAllComplete(listOfNotNull(poseDetectionTask)).addOnCompleteListener {
        onComplete(
            AnalysisResult(
                skeleton = skeleton,
                objects = objectsDetected
            ),
            bitmap
        )
    }
}