package com.example.livingai.data.ml

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Rect
import com.example.livingai.domain.ml.AIModel
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.common.FileUtil
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import java.nio.ByteBuffer
import java.nio.ByteOrder

class AIModelImpl(private val context: Context) : AIModel {

    private val objectDetector: Interpreter
    private val labels: List<String>

    init {
        // Load the TFLite model from assets
        val modelBuffer = FileUtil.loadMappedFile(context, "efficientdet-lite0.tflite")
        val options = Interpreter.Options().apply { numThreads = 4 }
        objectDetector = Interpreter(modelBuffer, options)

        // Load labels from assets
        labels = try {
            FileUtil.loadLabels(context, "labels.txt")
        } catch (e: Exception) {
            e.printStackTrace()
            emptyList()
        }
    }

    override suspend fun detectObject(bitmap: Bitmap): ObjectDetectionResult? {
        // Preprocess the image
        val imageProcessor = ImageProcessor.Builder()
            .add(ResizeOp(320, 320, ResizeOp.ResizeMethod.BILINEAR))
            .build()
        
        var tensorImage = TensorImage.fromBitmap(bitmap)
        tensorImage = imageProcessor.process(tensorImage)

        // Prepare model inputs and outputs
        // Based on crash: [1, 25, 4] vs [1, 10, 4]. The model outputs 25 detections, not 10.
        val locations = Array(1) { Array(25) { FloatArray(4) } }
        val classes = Array(1) { FloatArray(25) }
        val scores = Array(1) { FloatArray(25) }
        val numDetections = FloatArray(1)

        val outputs = mapOf(
            0 to locations,
            1 to classes,
            2 to scores,
            3 to numDetections
        )

        // Run inference
        objectDetector.runForMultipleInputsOutputs(arrayOf(tensorImage.buffer), outputs)

        // Post-process the results
        val bestDetection = scores[0].withIndex()
            .maxByOrNull { it.value }
            ?.takeIf { it.value > 0.5f } // Confidence threshold

        if (bestDetection != null) {
            val index = bestDetection.index
            val score = bestDetection.value
            val location = locations[0][index] // [ymin, xmin, ymax, xmax]
            val labelIndex = classes[0][index].toInt()
            val label = labels.getOrElse(labelIndex) { "Unknown" }

            // Convert normalized coordinates to absolute pixel values
            val ymin = location[0] * bitmap.height
            val xmin = location[1] * bitmap.width
            val ymax = location[2] * bitmap.height
            val xmax = location[3] * bitmap.width

            val boundingBox = Rect(xmin.toInt(), ymin.toInt(), xmax.toInt(), ymax.toInt())

            return ObjectDetectionResult(boundingBox, label, score)
        }

        return null
    }
    
    // This is no longer the primary function, but kept for interface compliance
    override suspend fun segmentImage(bitmap: Bitmap): Triple<Bitmap, BooleanArray, Rect>? {
        // Returning null as we are focusing on object detection now
        return null
    }

    override fun deriveInference(bitmap: Bitmap): String = "Object Detection"
}
