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 android.provider.OpenableColumns
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) {

    suspend fun segmentAndSave(inputUri: Uri): Uri? {
        return suspendCancellableCoroutine { continuation ->
            try {
                val image = InputImage.fromFilePath(context, inputUri)
                val options = SubjectSegmenterOptions.Builder()
                    .enableForegroundBitmap()
                    .build()
                val segmenter = SubjectSegmentation.getClient(options)

                segmenter.process(image)
                    .addOnSuccessListener { result ->
                        val foreground = result.foregroundBitmap
                        if (foreground != null) {
                            try {
                                val resultBitmap = Bitmap.createBitmap(
                                    foreground.width,
                                    foreground.height,
                                    Bitmap.Config.ARGB_8888
                                )
                                val canvas = Canvas(resultBitmap)
                                canvas.drawColor(Color.BLACK)
                                canvas.drawBitmap(foreground, 0f, 0f, null)

                                val originalName = getFileName(inputUri) ?: "image"
                                val nameWithoutExt = originalName.substringBeforeLast('.')
                                val baseName = nameWithoutExt.replace(Regex("_segmented.*"), "")
                                val filename = "${baseName}_segmented_${System.currentTimeMillis()}.jpg"

                                val contentValues = ContentValues().apply {
                                    put(MediaStore.MediaColumns.DISPLAY_NAME, filename)
                                    put(MediaStore.MediaColumns.MIME_TYPE, "image/jpeg")
                                    if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
                                        put(MediaStore.MediaColumns.RELATIVE_PATH, "Pictures/LivingAI/Segmented")
                                    }
                                }

                                val uri = context.contentResolver.insert(MediaStore.Images.Media.EXTERNAL_CONTENT_URI, contentValues)
                                
                                if (uri != null) {
                                    val outputStream: OutputStream? = context.contentResolver.openOutputStream(uri)
                                    outputStream?.use { out ->
                                        resultBitmap.compress(Bitmap.CompressFormat.JPEG, 100, out)
                                    }
                                    continuation.resume(uri)
                                } else {
                                    continuation.resume(null)
                                }
                            } catch (e: Exception) {
                                continuation.resumeWithException(e)
                            }
                        } else {
                            continuation.resume(null)
                        }
                    }
                    .addOnFailureListener { e ->
                        continuation.resumeWithException(e)
                    }
                    .addOnCompleteListener {
                        segmenter.close()
                    }
            } catch (e: Exception) {
                continuation.resumeWithException(e)
            }
        }
    }

    private fun getFileName(uri: Uri): String? {
        var result: String? = null
        if (uri.scheme == "content") {
            val cursor = context.contentResolver.query(uri, null, null, null, null)
            try {
                if (cursor != null && cursor.moveToFirst()) {
                    val index = cursor.getColumnIndex(OpenableColumns.DISPLAY_NAME)
                    if (index >= 0) {
                        result = cursor.getString(index)
                    }
                }
            } catch (e: Exception) {
                // ignore
            } finally {
                cursor?.close()
            }
        }
        if (result == null) {
            result = uri.path
            val cut = result?.lastIndexOf('/')
            if (cut != null && cut != -1) {
                result = result?.substring(cut + 1)
            }
        }
        return result
    }
}
