package com.example.livingai.domain.ml

import android.content.ContentValues
import android.content.Context
import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Color
import android.net.Uri
import android.os.Build
import android.provider.MediaStore
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.segmentation.subject.SubjectSegmentation
import com.google.mlkit.vision.segmentation.subject.SubjectSegmenterOptions
import kotlinx.coroutines.suspendCancellableCoroutine
import java.io.OutputStream
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

class SubjectSegmenterHelper(private val context: Context) {

    private suspend fun segmentInternal(image: InputImage): Bitmap? =
        suspendCancellableCoroutine { continuation ->

            val options = SubjectSegmenterOptions.Builder()
                .enableMultipleSubjects(
                    SubjectSegmenterOptions.SubjectResultOptions.Builder()
                        .enableSubjectBitmap()
                        .build()
                )
                .build()

            val segmenter = SubjectSegmentation.getClient(options)

            segmenter.process(image)
                .addOnSuccessListener { result ->
                    val subject = result.subjects
                        .maxByOrNull { it.width * it.height }

                    if (subject?.bitmap == null) {
                        continuation.resume(null)
                        return@addOnSuccessListener
                    }

                    try {
                        val output = Bitmap.createBitmap(
                            image.width,
                            image.height,
                            Bitmap.Config.ARGB_8888
                        )

                        val canvas = Canvas(output)
                        canvas.drawColor(Color.BLACK)
                        canvas.drawBitmap(
                            subject.bitmap!!,
                            subject.startX.toFloat(),
                            subject.startY.toFloat(),
                            null
                        )

                        continuation.resume(output)
                    } catch (e: Exception) {
                        continuation.resumeWithException(e)
                    }
                }
                .addOnFailureListener { e ->
                    continuation.resumeWithException(e)
                }
                .addOnCompleteListener {
                    segmenter.close()
                }
        }

    suspend fun segmentToBitmap(inputBitmap: Bitmap): Bitmap? {
        val image = InputImage.fromBitmap(inputBitmap, 0)
        return segmentInternal(image)
    }

    suspend fun segmentAndSave(
        inputBitmap: Bitmap,
        animalId: String,
        orientation: String,
        subFolder: String? = null
    ): Uri? {
        val image = InputImage.fromBitmap(inputBitmap, 0)
        val bitmap = segmentInternal(image) ?: return null
        return saveBitmap(bitmap, animalId, orientation, subFolder)
    }

    suspend fun segmentAndSave(
        inputUri: Uri,
        animalId: String,
        orientation: String,
        subFolder: String? = null
    ): Uri? {
        val image = InputImage.fromFilePath(context, inputUri)
        val bitmap = segmentInternal(image) ?: return null
        return saveBitmap(bitmap, animalId, orientation, subFolder)
    }

    private fun saveBitmap(
        bitmap: Bitmap,
        animalId: String,
        orientation: String,
        subFolder: String?
    ): Uri? {

        val filename = "${animalId}_${orientation}_segmented.jpg"

        val values = ContentValues().apply {
            put(MediaStore.MediaColumns.DISPLAY_NAME, filename)
            put(MediaStore.MediaColumns.MIME_TYPE, "image/jpeg")
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
                val path =
                    if (subFolder != null)
                        "Pictures/LivingAI/$animalId/$subFolder"
                    else
                        "Pictures/LivingAI/$animalId"
                put(MediaStore.MediaColumns.RELATIVE_PATH, path)
            }
        }

        val uri = context.contentResolver.insert(
            MediaStore.Images.Media.EXTERNAL_CONTENT_URI,
            values
        ) ?: return null

        val outputStream: OutputStream? =
            context.contentResolver.openOutputStream(uri)

        outputStream?.use { out ->
            bitmap.compress(Bitmap.CompressFormat.JPEG, 100, out)
        }

        return uri
    }
}
