package com.example.livingai.analysis

import android.content.res.Resources
import android.graphics.Bitmap
import android.graphics.Color
import android.graphics.Matrix
import android.util.Log
import com.example.livingai.commons.Constants
import com.google.android.gms.tasks.Task
import com.google.android.gms.tasks.TaskCompletionSource
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.segmentation.subject.SubjectSegmentation
import com.google.mlkit.vision.segmentation.subject.SubjectSegmenterOptions
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.math.max
import kotlin.math.min
import kotlin.math.sqrt

data class SegmentationResult(
    val mask: Bitmap?,
    val isMatch: Boolean
)

class SegmentProcessor {

    private val isProcessing = AtomicBoolean(false)
    private val processingExecutor = Executors.newSingleThreadExecutor()

    private val metrics = Resources.getSystem().displayMetrics
    private val screenW = metrics.widthPixels
    private val screenH = metrics.heightPixels

    private val options = SubjectSegmenterOptions.Builder()
        .enableMultipleSubjects(
            SubjectSegmenterOptions.SubjectResultOptions.Builder()
                .enableConfidenceMask()
                .build()
        )
        .build()

    private val segmenter = SubjectSegmentation.getClient(options)

    fun process(
        inputImage: InputImage,
        savedMask: Bitmap?,
        thresholdPercent: Int = 75,
        algorithm: String = Constants.ALGORITHM_HAMMING
    ): Task<SegmentationResult> {

        val tcs = TaskCompletionSource<SegmentationResult>()

        if (!isProcessing.compareAndSet(false, true)) {
            tcs.setResult(SegmentationResult(null, false))
            return tcs.task
        }

        segmenter.process(inputImage)
            .addOnSuccessListener(processingExecutor) { result ->

                var bitmapMask: Bitmap? = null
                val subject = result.subjects.firstOrNull()
                val mask = subject?.confidenceMask

                if (mask != null) {

                    val startX = subject.startX
                    val startY = subject.startY
                    val maskWidth = subject.width
                    val maskHeight = subject.height

                    val fullWidth = inputImage.width
                    val fullHeight = inputImage.height

                    if (mask.remaining() >= maskWidth * maskHeight) {

                        val colors = IntArray(fullWidth * fullHeight)
                        mask.rewind()

                        // Build full-resolution mask
                        for (y in 0 until maskHeight) {
                            for (x in 0 until maskWidth) {
                                if (mask.get() > 0.5f) {
                                    val dx = startX + x
                                    val dy = startY + y
                                    if (dx < fullWidth && dy < fullHeight) {
                                        colors[dy * fullWidth + dx] =
                                            Color.argb(180, 255, 0, 255)
                                    }
                                }
                            }
                        }

                        val rawBitmap = Bitmap.createBitmap(
                            colors, fullWidth, fullHeight, Bitmap.Config.ARGB_8888
                        )

                        // Rotate to match screen orientation if needed (so segmentation mask is upright)
                        if (inputImage.rotationDegrees != 0) {
                            val matrix = Matrix()
                            matrix.postRotate(inputImage.rotationDegrees.toFloat())
                            bitmapMask = Bitmap.createBitmap(
                                rawBitmap, 0, 0, rawBitmap.width, rawBitmap.height, matrix, true
                            )
                        } else {
                            bitmapMask = rawBitmap
                        }
                    }
                }

                var isMatch = false

                if (bitmapMask != null && savedMask != null) {
                    //To consider all kinds of phones with aspect ratio and screen size
                    val (scaledSavedMask, scaledSegmentedMask) = scaleBothToScreenAndCrop(savedMask, bitmapMask)

                    isMatch = when (algorithm) {
                        Constants.ALGORITHM_EUCLIDEAN ->
                            calculateEuclideanDistance(scaledSavedMask, scaledSegmentedMask, thresholdPercent)
                        Constants.ALGORITHM_JACCARD ->
                            calculateJaccardSimilarity(scaledSavedMask, scaledSegmentedMask, thresholdPercent)
                        else ->
                            calculateHammingDistance(scaledSavedMask, scaledSegmentedMask, thresholdPercent)
                    }
                }

                tcs.setResult(SegmentationResult(bitmapMask, isMatch))
            }
            .addOnFailureListener {
                Log.e("SegmentProcessor", "Segmentation failed", it)
                tcs.setException(it)
            }
            .addOnCompleteListener {
                isProcessing.set(false)
            }

        return tcs.task
    }

    private fun calculateHammingDistance(mask1: Bitmap, mask2: Bitmap, thresholdPercent: Int): Boolean {
        if (mask1.width != mask2.width || mask1.height != mask2.height) return false

        val width = mask1.width
        val height = mask1.height
        val p1 = IntArray(width * height)
        val p2 = IntArray(width * height)

        mask1.getPixels(p1, 0, width, 0, 0, width, height)
        mask2.getPixels(p2, 0, width, 0, 0, width, height)

        var distance = 0

        for (i in p1.indices) {
            val a = (p1[i] ushr 24) > 0
            val b = (p2[i] ushr 24) > 0
            if (a != b) distance++
        }

        val total = width * height
        val allowed = total * (100 - thresholdPercent) / 100

        return distance <= allowed
    }

    private fun calculateEuclideanDistance(mask1: Bitmap, mask2: Bitmap, thresholdPercent: Int): Boolean {
        if (mask1.width != mask2.width || mask1.height != mask2.height) return false

        val width = mask1.width
        val height = mask1.height
        val p1 = IntArray(width * height)
        val p2 = IntArray(width * height)

        mask1.getPixels(p1, 0, width, 0, 0, width, height)
        mask2.getPixels(p2, 0, width, 0, 0, width, height)

        var sum = 0L
        for (i in p1.indices) {
            val v1 = if ((p1[i] ushr 24) > 0) 255 else 0
            val v2 = if ((p2[i] ushr 24) > 0) 255 else 0
            val diff = v1 - v2
            sum += diff * diff
        }

        val dist = sqrt(sum.toDouble())
        val max = sqrt((width * height).toDouble()) * 255.0
        val allowed = max * (100 - thresholdPercent) / 100.0

        return dist <= allowed
    }

    private fun calculateJaccardSimilarity(mask1: Bitmap, mask2: Bitmap, thresholdPercent: Int): Boolean {
        if (mask1.width != mask2.width || mask1.height != mask2.height) return false

        val width = mask1.width
        val height = mask1.height
        val p1 = IntArray(width * height)
        val p2 = IntArray(width * height)

        mask1.getPixels(p1, 0, width, 0, 0, width, height)
        mask2.getPixels(p2, 0, width, 0, 0, width, height)

        var intersection = 0
        var union = 0

        for (i in p1.indices) {
            val a = ((p1[i] ushr 24) and 0xFF) > 0
            val b = ((p2[i] ushr 24) and 0xFF) > 0

            if (a && b) intersection++
            if (a || b) union++
        }

        if (union == 0) return false

        val score = (intersection.toDouble() / union.toDouble()) * 100.0

        return score >= thresholdPercent
    }

    fun centerCrop(bmp: Bitmap, tw: Int, th: Int): Bitmap {
        val offsetX = (bmp.width - tw) / 2
        val offsetY = (bmp.height - th) / 2
        return Bitmap.createBitmap(bmp, offsetX, offsetY, tw, th)
    }

    fun scaleFitCenter(bmp: Bitmap, isFit: Boolean): Bitmap {
        val w = bmp.width
        val h = bmp.height
        var scale: Float? = 0F
        // scale to fit inside screen
        if (isFit) {
            scale = min(screenW.toFloat() / w, screenH.toFloat() / h)
        } else {
            scale = max(screenW.toFloat() / w, screenH.toFloat() / h)
        }
        val newW = (w * scale).toInt()
        val newH = (h * scale).toInt()

        return Bitmap.createScaledBitmap(bmp, newW, newH, true)
    }

    fun scaleBothToScreenAndCrop(
        savedMask: Bitmap,
        segmentedMask: Bitmap
    ): Pair<Bitmap, Bitmap> {

        val scaledA = scaleFitCenter(savedMask, true)
        val scaledB = scaleFitCenter(segmentedMask, false)

        val wA = scaledA.width
        val hA = scaledA.height
        val wB = scaledB.width
        val hB = scaledB.height

        val targetW = min(wA, wB)
        val targetH = min(hA, hB)

        val croppedA = centerCrop(scaledA, targetW, targetH)
        val croppedB = centerCrop(scaledB, targetW, targetH)

        return Pair(croppedA, croppedB)
    }


}
