package com.performancecoachlab.posedetection.custom

import androidx.compose.ui.geometry.Rect
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.toPixelMap
import com.performancecoachlab.posedetection.recording.AnalysisObject
import com.performancecoachlab.posedetection.recording.FrameSize
import com.performancecoachlab.posedetection.recording.Label
import kotlinx.cinterop.CPointer
import kotlinx.cinterop.ExperimentalForeignApi
import kotlinx.cinterop.FloatVar
import kotlinx.cinterop.addressOf
import kotlinx.cinterop.get
import kotlinx.cinterop.reinterpret
import kotlinx.cinterop.usePinned
import platform.CoreGraphics.*
import platform.CoreML.MLMultiArray
import platform.Foundation.NSNumber
import platform.Vision.VNCoreMLFeatureValueObservation
import platform.Vision.VNCoreMLRequest
import platform.Vision.VNImageCropAndScaleOptionScaleFit
import platform.Vision.VNImageRequestHandler
import platform.Vision.VNRecognizedObjectObservation
import kotlin.math.max
import kotlin.math.min

@OptIn(ExperimentalForeignApi::class)
actual class ImageDetector actual constructor(model: ObjectModel) {

    private val vncoreModel = model.getModel()
    private val inputW = model.inputWidth
    private val inputH = model.inputHeight

    actual fun detect(image: ImageBitmap): List<AnalysisObject> {
        val model = vncoreModel ?: return emptyList()
        val imgW = image.width
        val imgH = image.height
        if (imgW <= 0 || imgH <= 0) return emptyList()

        val srcCgImage = createCgImageFromImageBitmap(image) ?: return emptyList()

        // If we know the model's input size, letterbox the source into
        // (inputW × inputH) with gray-114 padding — matching ultralytics'
        // training-time preprocessing. Otherwise fall back to passing the
        // raw image and letting Vision scale it.
        val useLetterbox = inputW > 0 && inputH > 0
        val scale: Float
        val padX: Float
        val padY: Float
        val handlerImage: CPointer<cnames.structs.CGImage>
        if (useLetterbox) {
            val s = min(
                inputW.toFloat() / imgW.toFloat(),
                inputH.toFloat() / imgH.toFloat()
            )
            val scaledW = (imgW * s)
            val scaledH = (imgH * s)
            val px = (inputW - scaledW) / 2f
            val py = (inputH - scaledH) / 2f
            val letterboxed = createLetterboxedCgImage(srcCgImage, inputW, inputH, px, py, scaledW, scaledH)
            CGImageRelease(srcCgImage)
            if (letterboxed == null) return emptyList()
            scale = s
            padX = px
            padY = py
            handlerImage = letterboxed
        } else {
            scale = 1f
            padX = 0f
            padY = 0f
            handlerImage = srcCgImage
        }

        val results = mutableListOf<AnalysisObject>()
        val bbW = if (useLetterbox) inputW.toFloat() else imgW.toFloat()
        val bbH = if (useLetterbox) inputH.toFloat() else imgH.toFloat()
        val request = VNCoreMLRequest(model) { req, _ ->
            val observations = req?.results as? List<*> ?: return@VNCoreMLRequest
            for (obs in observations) {
                when (obs) {
                    is VNRecognizedObjectObservation -> {
                        // Vision-compatible model (classifier + coordinates pipeline).
                        // Coords: origin bottom-left, normalized 0-1 over the handler image.
                        val label = obs.labels.firstOrNull()?.toString() ?: "Unknown"
                        val confidence = obs.confidence
                        val bb = obs.boundingBox
                        val x1LB = CGRectGetMinX(bb).toFloat() * bbW
                        val y1LB = (1f - CGRectGetMaxY(bb).toFloat()) * bbH
                        val x2LB = CGRectGetMaxX(bb).toFloat() * bbW
                        val y2LB = (1f - CGRectGetMinY(bb).toFloat()) * bbH
                        addDetection(
                            results, x1LB, y1LB, x2LB, y2LB, padX, padY, scale,
                            imgW, imgH, label, confidence
                        )
                    }
                    is VNCoreMLFeatureValueObservation -> {
                        // Raw multiarray output (e.g. yolo26n end2end). Shape
                        // [1, 300, 6] with rows [x1, y1, x2, y2, conf, cls] in
                        // normalized top-left-origin coords over the handler image.
                        val arr = obs.featureValue.multiArrayValue ?: continue
                        decodeRawDetections(
                            arr, results, bbW, bbH, padX, padY, scale, imgW, imgH
                        )
                    }
                }
            }
        }.apply {
            // Belt-and-suspenders: input already matches model aspect when
            // letterboxed, so this is effectively a no-op — but prevents
            // Vision from re-cropping if anything upstream shifts.
            imageCropAndScaleOption = VNImageCropAndScaleOptionScaleFit
        }

        val handler = VNImageRequestHandler(handlerImage, mapOf<Any?, Any?>())
        handler.performRequests(listOf(request), null)

        CGImageRelease(handlerImage)
        return results
    }
}

private fun addDetection(
    results: MutableList<AnalysisObject>,
    x1LB: Float, y1LB: Float, x2LB: Float, y2LB: Float,
    padX: Float, padY: Float, scale: Float,
    imgW: Int, imgH: Int,
    label: String, confidence: Float,
) {
    val left = ((x1LB - padX) / scale).coerceIn(0f, imgW.toFloat())
    val top = ((y1LB - padY) / scale).coerceIn(0f, imgH.toFloat())
    val right = ((x2LB - padX) / scale).coerceIn(0f, imgW.toFloat())
    val bottom = ((y2LB - padY) / scale).coerceIn(0f, imgH.toFloat())
    results.add(
        AnalysisObject(
            boundingBox = Rect(left, top, right, bottom),
            trackingId = 0,
            labels = listOf(Label(label, confidence)),
            frameSize = FrameSize(imgW, imgH),
            timestamp = 0L
        )
    )
}

private val CLASS_NAMES = listOf("basketball", "basketball_hoop")

@OptIn(ExperimentalForeignApi::class)
private fun decodeRawDetections(
    arr: MLMultiArray,
    results: MutableList<AnalysisObject>,
    bbW: Float, bbH: Float,
    padX: Float, padY: Float, scale: Float,
    imgW: Int, imgH: Int,
) {
    val shape = arr.shape
    if (shape.size != 3) return
    val dim1 = (shape[1] as NSNumber).intValue
    val dim2 = (shape[2] as NSNumber).intValue
    // Expect [1, N, 6] (yolo end2end).
    if (dim2 != 6) return
    val elements = dim1
    val dataPtr = arr.dataPointer?.reinterpret<FloatVar>() ?: return
    fun at(i: Int, j: Int): Float = dataPtr[i * 6 + j]
    for (i in 0 until elements) {
        val conf = at(i, 4)
        if (conf <= 0.25f) continue
        val x1n = at(i, 0)
        val y1n = at(i, 1)
        val x2n = at(i, 2)
        val y2n = at(i, 3)
        val cls = at(i, 5).toInt()
        // Top-left-origin normalized coords over the handler image.
        val x1LB = min(x1n, x2n) * bbW
        val y1LB = min(y1n, y2n) * bbH
        val x2LB = max(x1n, x2n) * bbW
        val y2LB = max(y1n, y2n) * bbH
        val label = CLASS_NAMES.getOrNull(cls) ?: "class_$cls"
        addDetection(
            results, x1LB, y1LB, x2LB, y2LB, padX, padY, scale,
            imgW, imgH, label, conf
        )
    }
}

@OptIn(ExperimentalForeignApi::class)
private fun createCgImageFromImageBitmap(image: ImageBitmap): CPointer<cnames.structs.CGImage>? {
    val w = image.width
    val h = image.height
    val pixelMap = image.toPixelMap()
    val buffer = ByteArray(w * h * 4)
    for (y in 0 until h) {
        for (x in 0 until w) {
            val color = pixelMap[x, y]
            val idx = (y * w + x) * 4
            buffer[idx] = (color.red * 255).toInt().toByte()
            buffer[idx + 1] = (color.green * 255).toInt().toByte()
            buffer[idx + 2] = (color.blue * 255).toInt().toByte()
            buffer[idx + 3] = (color.alpha * 255).toInt().toByte()
        }
    }
    val colorSpace = CGColorSpaceCreateDeviceRGB()
    val bitmapInfo = CGImageAlphaInfo.kCGImageAlphaPremultipliedLast.value
    val context = CGBitmapContextCreate(
        null, w.toULong(), h.toULong(), 8u,
        (w * 4).toULong(), colorSpace, bitmapInfo
    ) ?: run {
        CGColorSpaceRelease(colorSpace)
        return null
    }
    buffer.usePinned { pinned ->
        val data = CGBitmapContextGetData(context)
        if (data != null) {
            platform.posix.memcpy(data, pinned.addressOf(0), (w * h * 4).toULong())
        }
    }
    val cgImage = CGBitmapContextCreateImage(context)
    CGContextRelease(context)
    CGColorSpaceRelease(colorSpace)
    return cgImage
}

@OptIn(ExperimentalForeignApi::class)
private fun createLetterboxedCgImage(
    src: CPointer<cnames.structs.CGImage>,
    outW: Int,
    outH: Int,
    padX: Float,
    padY: Float,
    drawW: Float,
    drawH: Float,
): CPointer<cnames.structs.CGImage>? {
    val colorSpace = CGColorSpaceCreateDeviceRGB()
    val bitmapInfo = CGImageAlphaInfo.kCGImageAlphaPremultipliedLast.value
    val ctx = CGBitmapContextCreate(
        null, outW.toULong(), outH.toULong(), 8u,
        (outW * 4).toULong(), colorSpace, bitmapInfo
    ) ?: run {
        CGColorSpaceRelease(colorSpace)
        return null
    }

    // Gray-114 fill to match ultralytics letterbox padding.
    CGContextSetRGBFillColor(ctx, 114.0 / 255.0, 114.0 / 255.0, 114.0 / 255.0, 1.0)
    CGContextFillRect(ctx, CGRectMake(0.0, 0.0, outW.toDouble(), outH.toDouble()))

    // Core Graphics has bottom-left origin; our padY/drawH are top-left,
    // so flip the Y coordinate before drawing.
    val flippedY = outH.toFloat() - padY - drawH
    CGContextDrawImage(
        ctx,
        CGRectMake(padX.toDouble(), flippedY.toDouble(), drawW.toDouble(), drawH.toDouble()),
        src
    )

    val image = CGBitmapContextCreateImage(ctx)
    CGContextRelease(ctx)
    CGColorSpaceRelease(colorSpace)
    return image
}
