package com.performancecoachlab.posedetection.recording

import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Paint
import androidx.compose.ui.geometry.Rect
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asImageBitmap
import androidx.core.graphics.createBitmap
import co.touchlab.kermit.Logger
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.pose.PoseDetection
import com.google.mlkit.vision.pose.accurate.AccuratePoseDetectorOptions
import com.performancecoachlab.posedetection.camera.applyFocusAreaMaskPooled
import com.performancecoachlab.posedetection.camera.drawAnalysisResults
import com.performancecoachlab.posedetection.camera.drawSkeleton
import com.performancecoachlab.posedetection.camera.process
import com.performancecoachlab.posedetection.custom.ObjectModel
import com.performancecoachlab.posedetection.skeleton.Skeleton
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import org.tensorflow.lite.DataType
import org.tensorflow.lite.support.common.ops.CastOp
import org.tensorflow.lite.support.common.ops.NormalizeOp
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import kotlin.coroutines.resume
import kotlin.math.max

actual class InputFrame(val bitmap: Bitmap, actual val timestamp: Long) {
    actual fun toImageBitmap(): ImageBitmap {
        return bitmap.asImageBitmap()
    }

    actual fun drawSkeleton(skel: Skeleton): ImageBitmap {
        return toImageBitmap().drawSkeleton(skel)
    }

    actual fun drawAnalysisResults(analysisResults: AnalysisResult): ImageBitmap {
        return toImageBitmap().drawAnalysisResults(analysisResults)
    }

    actual fun crop(rect: Rect): InputFrame {
        val width = bitmap.width
        val height = bitmap.height
        val sanitised = Rect(
            left = rect.left.coerceIn(0f, width.toFloat()),
            top = rect.top.coerceIn(0f, height.toFloat()),
            right = rect.right.coerceIn(0f, width.toFloat()),
            bottom = rect.bottom.coerceIn(0f, height.toFloat())
        )
        val croppedBitmap = Bitmap.createBitmap(
            bitmap,
            sanitised.left.toInt(),
            sanitised.top.toInt(),
            sanitised.width.toInt(),
            sanitised.height.toInt()
        )
        return InputFrame(croppedBitmap, timestamp)
    }
}

actual class FrameAnalyser actual constructor(val model: ObjectModel?) {
    private val options =
        AccuratePoseDetectorOptions.Builder().setDetectorMode(AccuratePoseDetectorOptions.STREAM_MODE).build()
    private val poseDetector = PoseDetection.getClient(options)
    private val objDetector = model?.getDetector()

    private val imageProcessor = ImageProcessor.Builder()
        .add(NormalizeOp(0f, 255f))
        .add(CastOp(DataType.FLOAT32))
        .build()

    private var tensorBitmap: Bitmap? = null
    private var tensorW: Int = 0
    private var tensorH: Int = 0

    private fun ensureTensorBitmap(width: Int, height: Int): Bitmap {
        val existing = tensorBitmap
        return if (existing != null && !existing.isRecycled && tensorW == width && tensorH == height) {
            existing
        } else {
            androidx.core.graphics.createBitmap(width, height, Bitmap.Config.ARGB_8888).also {
                tensorBitmap = it
                tensorW = width
                tensorH = height
            }
        }
    }

    private fun toPoseBitmap(src: Bitmap): Bitmap {
        val argb8888 = if (src.config == Bitmap.Config.ARGB_8888) src else src.copy(
            Bitmap.Config.ARGB_8888,
            false
        )

        val minSide = 480
        val w = argb8888.width
        val h = argb8888.height
        val scale = max(minSide.toFloat() / w.toFloat(), minSide.toFloat() / h.toFloat())
        if (scale <= 1f) return argb8888

        val outW = (w * scale).toInt().coerceAtLeast(1)
        val outH = (h * scale).toInt().coerceAtLeast(1)
        val resized = createBitmap(outW, outH)
        Canvas(resized).drawBitmap(
            argb8888,
            null,
            android.graphics.Rect(0, 0, outW, outH),
            Paint(Paint.FILTER_BITMAP_FLAG)
        )
        return resized
    }

    actual suspend fun analyseFrame(inputFrame: InputFrame, focusArea: Rect?): AnalysisResult =
        suspendCancellableCoroutine { continuation ->
            val interpreter = objDetector

            // Pose: keep original resolution (optionally masked) for ML Kit.
            val masked = inputFrame.bitmap.applyFocusAreaMaskPooled(focusArea, 0)
            val mlKitImage = InputImage.fromBitmap(masked, 0)

            // Object: resize to model input size, then normalize.
            val tensorImage: TensorImage? = interpreter?.modelInfo?.let { tfl ->
                val w = tfl.inputWidth
                val h = tfl.inputHeight

                if (w <= 0 || h <= 0) null else {
                    val dst = ensureTensorBitmap(w, h)
                    Canvas(dst).drawBitmap(
                        inputFrame.bitmap,
                        null,
                        android.graphics.Rect(0, 0, w, h),
                        android.graphics.Paint(android.graphics.Paint.FILTER_BITMAP_FLAG)
                    )
                    TensorImage(DataType.FLOAT32).also { ti -> ti.load(dst) }
                        .let(imageProcessor::process)
                }
            }
            continuation.context.let { ctx ->
                kotlinx.coroutines.CoroutineScope(ctx).launch(Dispatchers.Default) {
                    try {
                        process(
                            tensorImage = tensorImage,
                            mlKitImage = mlKitImage,
                            objectDetector = interpreter,
                            poseDetector = poseDetector,
                            timestamp = inputFrame.timestamp,
                            width = inputFrame.bitmap.width,
                            height = inputFrame.bitmap.height,
                            bitmap = inputFrame.bitmap,
                            onComplete = { result, _ ->
                                if (continuation.isActive) continuation.resume(result)
                            }
                        )
                    } catch (t: Throwable) {
                        Logger.e(t) { "Offline analyseFrame: process() crashed" }
                        if (continuation.isActive) {
                            continuation.resume(
                                AnalysisResult(
                                    skeleton = null,
                                    objects = emptyList()
                                )
                            )
                        }
                    }
                }
            }
        }
}