package com.performancecoachlab.posedetection.custom

import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.graphics.RectF
import androidx.compose.ui.geometry.Rect
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asAndroidBitmap
import co.touchlab.kermit.Logger
import com.performancecoachlab.posedetection.camera.label
import com.performancecoachlab.posedetection.recording.AnalysisObject
import com.performancecoachlab.posedetection.recording.FrameSize
import com.performancecoachlab.posedetection.recording.Label
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 org.tensorflow.lite.support.tensorbuffer.TensorBuffer
import kotlin.math.absoluteValue
import kotlin.math.max
import kotlin.math.min
import kotlin.math.roundToInt

actual class ImageDetector actual constructor(model: ObjectModel) {

    private val detector: AndroidDetector? = model.getDetector()

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

    actual fun detect(image: ImageBitmap): List<AnalysisObject> {
        val det = detector ?: return emptyList()
        val info = det.modelInfo
        val inputW = info.inputWidth
        val inputH = info.inputHeight
        if (inputW <= 0 || inputH <= 0) return emptyList()

        val srcBitmap = image.asAndroidBitmap().let { bmp ->
            if (bmp.config == Bitmap.Config.ARGB_8888) bmp
            else bmp.copy(Bitmap.Config.ARGB_8888, false)
        }
        val imgW = srcBitmap.width
        val imgH = srcBitmap.height

        // Letterbox: scale the source to fit inside (inputW, inputH) while
        // preserving aspect ratio, center it, and pad the remainder with
        // gray 114 — matching ultralytics training-time preprocessing.
        // (A naive stretch-resize destroys aspect ratio and hurts detection
        // quality on non-square camera frames.)
        val scale = min(
            inputW.toFloat() / imgW.toFloat(),
            inputH.toFloat() / imgH.toFloat()
        )
        val scaledW = (imgW * scale).roundToInt().coerceAtLeast(1)
        val scaledH = (imgH * scale).roundToInt().coerceAtLeast(1)
        val padX = (inputW - scaledW) / 2f
        val padY = (inputH - scaledH) / 2f

        val resized = Bitmap.createBitmap(inputW, inputH, Bitmap.Config.ARGB_8888)
        val canvas = Canvas(resized)
        canvas.drawColor(Color.rgb(114, 114, 114))
        canvas.drawBitmap(
            srcBitmap, null,
            RectF(padX, padY, padX + scaledW, padY + scaledH),
            Paint(Paint.FILTER_BITMAP_FLAG)
        )

        // Normalize and convert to tensor
        val tensorImage = TensorImage(DataType.FLOAT32).also { it.load(resized) }
            .let(imageProcessor::process)

        // Run inference
        val outputShape = info.outputShape
        val output = TensorBuffer.createFixedSize(outputShape, DataType.FLOAT32)
        try {
            det.interpreter.run(tensorImage.buffer, output.buffer)
        } catch (e: Exception) {
            Logger.e(e) { "ImageDetector: interpreter.run failed" }
            return emptyList()
        }

        // Parse output
        val array = output.floatArray
        if (outputShape.size != 3) return emptyList()

        val dim1 = outputShape[1]
        val dim2 = outputShape[2]

        val elements: Int
        val channels: Int
        val isElementsFirst: Boolean

        when {
            dim2 == 6 -> { elements = dim1; channels = dim2; isElementsFirst = true }
            dim1 == 6 -> { channels = dim1; elements = dim2; isElementsFirst = false }
            else -> return emptyList()
        }

        fun valueAt(elementIndex: Int, channelIndex: Int): Float {
            return if (isElementsFirst) array[elementIndex * channels + channelIndex]
            else array[channelIndex * elements + elementIndex]
        }

        val imgWF = imgW.toFloat()
        val imgHF = imgH.toFloat()

        return (0 until elements).mapNotNull { i ->
            val cnf = valueAt(i, 4)
            if (cnf > 0.25f) {
                val x1 = valueAt(i, 0)
                val y1 = valueAt(i, 1)
                val x2 = valueAt(i, 2)
                val y2 = valueAt(i, 3)
                val cls = valueAt(i, 5).toInt()

                // Model outputs are normalized [0,1] over the letterboxed
                // input (inputW × inputH). Un-letterbox: convert to
                // letterboxed pixel coords, subtract padding, divide by
                // the fit ratio to get original-image pixel coords.
                val x1pLb = min(x1, x2) * inputW
                val y1pLb = min(y1, y2) * inputH
                val x2pLb = max(x1, x2) * inputW
                val y2pLb = max(y1, y2) * inputH

                val left = ((x1pLb - padX) / scale).coerceIn(0f, imgWF)
                val top = ((y1pLb - padY) / scale).coerceIn(0f, imgHF)
                val right = ((x2pLb - padX) / scale).coerceIn(0f, imgWF)
                val bottom = ((y2pLb - padY) / scale).coerceIn(0f, imgHF)

                AnalysisObject(
                    boundingBox = Rect(
                        left = left,
                        top = top,
                        right = right,
                        bottom = bottom
                    ),
                    trackingId = 0,
                    labels = listOf(Label(info.label(cls), cnf)),
                    frameSize = FrameSize(width = imgW.absoluteValue, height = imgH.absoluteValue),
                    timestamp = 0L
                )
            } else null
        }
    }
}
