package com.example.livingai.domain.ml

import android.content.Context
import android.graphics.Bitmap
import android.graphics.RectF
import org.tensorflow.lite.DataType
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import org.tensorflow.lite.support.image.ops.Rot90Op
import java.io.FileInputStream
import java.nio.channels.FileChannel

class ObjectDetectorImpl(
    private val context: Context,
    override var onResults: (List<ObjectDetector.DetectionResult>, Long) -> Unit,
    override var onError: (String) -> Unit
) : ObjectDetector {

    private var interpreter: Interpreter? = null
    private val modelName = "efficientdet-lite0.tflite"
    private val inputSize = 320 // EfficientDet-Lite0 expects 320x320

    init {
        setupInterpreter()
    }

    private fun setupInterpreter() {
        try {
            val assetFileDescriptor = context.assets.openFd(modelName)
            val fileInputStream = FileInputStream(assetFileDescriptor.fileDescriptor)
            val fileChannel = fileInputStream.channel
            val startOffset = assetFileDescriptor.startOffset
            val declaredLength = assetFileDescriptor.declaredLength
            val modelBuffer = fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
            
            val options = Interpreter.Options()
            options.setNumThreads(4)
            interpreter = Interpreter(modelBuffer, options)
        } catch (e: Exception) {
            onError(e.message ?: "Error loading model")
        }
    }

    override fun detect(bitmap: Bitmap, imageRotation: Int) {
        val tflite = interpreter ?: return

        val startTime = System.currentTimeMillis()

        // 1. Preprocess Image
        // Rotate -> Resize -> Convert to TensorImage (UINT8)
        val imageProcessor = ImageProcessor.Builder()
            .add(Rot90Op(-imageRotation / 90))
            .add(ResizeOp(inputSize, inputSize, ResizeOp.ResizeMethod.BILINEAR))
            .build()

        var tensorImage = TensorImage(DataType.UINT8)
        tensorImage.load(bitmap)
        tensorImage = imageProcessor.process(tensorImage)

        // 2. Prepare Output Buffers
        val outputBoxes = Array(1) { Array(25) { FloatArray(4) } }
        val outputClasses = Array(1) { FloatArray(25) }
        val outputScores = Array(1) { FloatArray(25) }
        val outputCount = FloatArray(1)

        val outputs = mapOf(
            0 to outputBoxes,
            1 to outputClasses,
            2 to outputScores,
            3 to outputCount
        )

        // 3. Run Inference
        try {
            tflite.runForMultipleInputsOutputs(arrayOf(tensorImage.buffer), outputs)
        } catch (e: Exception) {
            onError(e.message ?: "Inference failed")
            return
        }

        val inferenceTime = System.currentTimeMillis() - startTime

        // 4. Parse Results
        val results = mutableListOf<ObjectDetector.DetectionResult>()
        
        // Calculate dimensions of the rotated image (the coordinate space of the detections)
        val rotatedWidth = if (imageRotation % 180 == 0) bitmap.width else bitmap.height
        val rotatedHeight = if (imageRotation % 180 == 0) bitmap.height else bitmap.width

        for (i in 0 until 25) {
            val score = outputScores[0][i]
            if (score < 0.4f) continue 

            val classId = outputClasses[0][i].toInt()
            val label = if (classId == 20) "Cow" else "Object $classId"
            
            // Get box: [ymin, xmin, ymax, xmax] (normalized 0..1)
            val box = outputBoxes[0][i]
            val ymin = box[0]
            val xmin = box[1]
            val ymax = box[2]
            val xmax = box[3]

            // Scale to rotated image dimensions
            val left = xmin * rotatedWidth
            val top = ymin * rotatedHeight
            val right = xmax * rotatedWidth
            val bottom = ymax * rotatedHeight

            val boundingBox = RectF(left, top, right, bottom)
            
            results.add(ObjectDetector.DetectionResult(boundingBox, label, score))
        }

        onResults(results, inferenceTime)
    }
}
