tilt-shift-camera/app/src/main/java/no/naiv/tiltshift/camera/ImageCaptureHandler.kt

package no.naiv.tiltshift.camera

import android.content.Context
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Matrix
import android.location.Location
import androidx.camera.core.ImageCapture
import androidx.camera.core.ImageCaptureException
import androidx.camera.core.ImageProxy
import androidx.exifinterface.media.ExifInterface
import kotlinx.coroutines.suspendCancellableCoroutine
import no.naiv.tiltshift.effect.BlurMode
import no.naiv.tiltshift.effect.BlurParameters
import no.naiv.tiltshift.storage.PhotoSaver
import no.naiv.tiltshift.storage.SaveResult
import java.util.concurrent.Executor
import kotlin.coroutines.resume
import kotlin.math.cos
import kotlin.math.sin
import kotlin.math.sqrt

/**
 * Handles capturing photos with the tilt-shift effect applied.
 */
class ImageCaptureHandler(
    private val context: Context,
    private val photoSaver: PhotoSaver
) {

    /**
     * Holds the processed bitmap ready for saving, produced inside the
     * camera callback (synchronous CPU work) and consumed afterwards
     * in the caller's coroutine context.
     */
    private class ProcessedCapture(val bitmap: Bitmap)

    /**
     * Captures a photo and applies the tilt-shift effect.
     *
     * Phase 1 (inside suspendCancellableCoroutine / camera callback):
     *   decode → rotate → apply effect  (synchronous CPU work only)
     *
     * Phase 2 (after continuation resumes, back in coroutine context):
     *   save bitmap via PhotoSaver (suspend-safe)
     */
    suspend fun capturePhoto(
        imageCapture: ImageCapture,
        executor: Executor,
        blurParams: BlurParameters,
        deviceRotation: Int,
        location: Location?,
        isFrontCamera: Boolean
    ): SaveResult {
        // Phase 1: capture and process the image synchronously in the callback
        val captureResult = suspendCancellableCoroutine { continuation ->
            imageCapture.takePicture(
                executor,
                object : ImageCapture.OnImageCapturedCallback() {
                    override fun onCaptureSuccess(imageProxy: ImageProxy) {
                        try {
                            val imageRotation = imageProxy.imageInfo.rotationDegrees

                            var bitmap = imageProxyToBitmap(imageProxy)
                            imageProxy.close()

                            if (bitmap == null) {
                                continuation.resume(
                                    SaveResult.Error("Failed to convert image") as Any
                                )
                                return
                            }

                            bitmap = rotateBitmap(bitmap, imageRotation, isFrontCamera)

                            val processedBitmap = applyTiltShiftEffect(bitmap, blurParams)
                            bitmap.recycle()

                            continuation.resume(ProcessedCapture(processedBitmap))
                        } catch (e: Exception) {
                            continuation.resume(
                                SaveResult.Error("Capture failed: ${e.message}", e) as Any
                            )
                        }
                    }

                    override fun onError(exception: ImageCaptureException) {
                        continuation.resume(
                            SaveResult.Error(
                                "Capture failed: ${exception.message}", exception
                            ) as Any
                        )
                    }
                }
            )
        }

        // Phase 2: save to disk in the caller's coroutine context (suspend-safe)
        if (captureResult is ProcessedCapture) {
            return try {
                photoSaver.saveBitmap(
                    captureResult.bitmap,
                    ExifInterface.ORIENTATION_NORMAL,
                    location
                )
            } finally {
                captureResult.bitmap.recycle()
            }
        }

        return captureResult as SaveResult
    }

    /**
     * Rotates a bitmap to the correct orientation.
     */
    private fun rotateBitmap(bitmap: Bitmap, rotationDegrees: Int, isFrontCamera: Boolean): Bitmap {
        if (rotationDegrees == 0 && !isFrontCamera) {
            return bitmap
        }

        val matrix = Matrix()

        // Apply rotation
        if (rotationDegrees != 0) {
            matrix.postRotate(rotationDegrees.toFloat())
        }

        // Mirror for front camera
        if (isFrontCamera) {
            matrix.postScale(-1f, 1f)
        }

        val rotated = Bitmap.createBitmap(
            bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true
        )

        if (rotated != bitmap) {
            bitmap.recycle()
        }

        return rotated
    }

    private fun imageProxyToBitmap(imageProxy: ImageProxy): Bitmap? {
        val buffer = imageProxy.planes[0].buffer
        val bytes = ByteArray(buffer.remaining())
        buffer.get(bytes)
        return BitmapFactory.decodeByteArray(bytes, 0, bytes.size)
    }

    /**
     * Applies tilt-shift blur effect to a bitmap.
     * Supports both linear and radial modes.
     */
    private fun applyTiltShiftEffect(source: Bitmap, params: BlurParameters): Bitmap {
        val width = source.width
        val height = source.height

        // Create output bitmap
        val result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)

        // For performance, we use a scaled-down version for blur and composite
        val scaleFactor = 4 // Blur a 1/4 size image for speed
        val blurredWidth = width / scaleFactor
        val blurredHeight = height / scaleFactor

        // Create scaled bitmap for blur
        val scaled = Bitmap.createScaledBitmap(source, blurredWidth, blurredHeight, true)

        // Apply stack blur (fast approximation)
        val blurred = stackBlur(scaled, (params.blurAmount * 25).toInt().coerceIn(1, 25))
        scaled.recycle()

        // Scale blurred back up
        val blurredFullSize = Bitmap.createScaledBitmap(blurred, width, height, true)
        blurred.recycle()

        // Create gradient mask based on tilt-shift parameters
        val mask = createGradientMask(width, height, params)

        // Composite: blend original with blurred based on mask
        val pixels = IntArray(width * height)
        val blurredPixels = IntArray(width * height)
        val maskPixels = IntArray(width * height)

        source.getPixels(pixels, 0, width, 0, 0, width, height)
        blurredFullSize.getPixels(blurredPixels, 0, width, 0, 0, width, height)
        mask.getPixels(maskPixels, 0, width, 0, 0, width, height)

        blurredFullSize.recycle()
        mask.recycle()

        for (i in pixels.indices) {
            val maskAlpha = (maskPixels[i] and 0xFF) / 255f
            val origR = (pixels[i] shr 16) and 0xFF
            val origG = (pixels[i] shr 8) and 0xFF
            val origB = pixels[i] and 0xFF
            val blurR = (blurredPixels[i] shr 16) and 0xFF
            val blurG = (blurredPixels[i] shr 8) and 0xFF
            val blurB = blurredPixels[i] and 0xFF

            val r = (origR * (1 - maskAlpha) + blurR * maskAlpha).toInt()
            val g = (origG * (1 - maskAlpha) + blurG * maskAlpha).toInt()
            val b = (origB * (1 - maskAlpha) + blurB * maskAlpha).toInt()

            pixels[i] = (0xFF shl 24) or (r shl 16) or (g shl 8) or b
        }

        result.setPixels(pixels, 0, width, 0, 0, width, height)
        return result
    }

    /**
     * Creates a gradient mask for the tilt-shift effect.
     * Supports both linear and radial modes.
     */
    private fun createGradientMask(width: Int, height: Int, params: BlurParameters): Bitmap {
        val mask = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
        val pixels = IntArray(width * height)

        val centerX = width * params.positionX
        val centerY = height * params.positionY
        val focusSize = height * params.size * 0.5f
        val transitionSize = focusSize * params.falloff

        val cosAngle = cos(params.angle)
        val sinAngle = sin(params.angle)
        val screenAspect = width.toFloat() / height.toFloat()

        for (y in 0 until height) {
            for (x in 0 until width) {
                val dist = when (params.mode) {
                    BlurMode.LINEAR -> {
                        // Rotate point around focus center
                        val dx = x - centerX
                        val dy = y - centerY
                        val rotatedY = -dx * sinAngle + dy * cosAngle
                        kotlin.math.abs(rotatedY)
                    }
                    BlurMode.RADIAL -> {
                        // Calculate elliptical distance from center
                        var dx = x - centerX
                        var dy = y - centerY

                        // Adjust for screen aspect ratio
                        dx *= screenAspect

                        // Rotate
                        val rotatedX = dx * cosAngle - dy * sinAngle
                        val rotatedY = dx * sinAngle + dy * cosAngle

                        // Apply ellipse aspect ratio
                        val adjustedX = rotatedX / params.aspectRatio

                        sqrt(adjustedX * adjustedX + rotatedY * rotatedY)
                    }
                }

                // Calculate blur amount based on distance from focus region
                val blurAmount = when {
                    dist < focusSize -> 0f
                    dist < focusSize + transitionSize -> {
                        (dist - focusSize) / transitionSize
                    }
                    else -> 1f
                }

                val gray = (blurAmount * 255).toInt().coerceIn(0, 255)
                pixels[y * width + x] = (0xFF shl 24) or (gray shl 16) or (gray shl 8) or gray
            }
        }

        mask.setPixels(pixels, 0, width, 0, 0, width, height)
        return mask
    }

    /**
     * Fast stack blur algorithm.
     */
    private fun stackBlur(bitmap: Bitmap, radius: Int): Bitmap {
        if (radius < 1) return bitmap.copy(Bitmap.Config.ARGB_8888, true)

        val w = bitmap.width
        val h = bitmap.height
        val pix = IntArray(w * h)
        bitmap.getPixels(pix, 0, w, 0, 0, w, h)

        val wm = w - 1
        val hm = h - 1
        val wh = w * h
        val div = radius + radius + 1

        val r = IntArray(wh)
        val g = IntArray(wh)
        val b = IntArray(wh)
        var rsum: Int
        var gsum: Int
        var bsum: Int
        var x: Int
        var y: Int
        var i: Int
        var p: Int
        var yp: Int
        var yi: Int
        var yw: Int
        val vmin = IntArray(maxOf(w, h))

        var divsum = (div + 1) shr 1
        divsum *= divsum
        val dv = IntArray(256 * divsum)
        for (i2 in 0 until 256 * divsum) {
            dv[i2] = (i2 / divsum)
        }

        yi = 0
        yw = 0

        val stack = Array(div) { IntArray(3) }
        var stackpointer: Int
        var stackstart: Int
        var sir: IntArray
        var rbs: Int
        val r1 = radius + 1
        var routsum: Int
        var goutsum: Int
        var boutsum: Int
        var rinsum: Int
        var ginsum: Int
        var binsum: Int

        for (y2 in 0 until h) {
            rinsum = 0
            ginsum = 0
            binsum = 0
            routsum = 0
            goutsum = 0
            boutsum = 0
            rsum = 0
            gsum = 0
            bsum = 0
            for (i2 in -radius..radius) {
                p = pix[yi + minOf(wm, maxOf(i2, 0))]
                sir = stack[i2 + radius]
                sir[0] = (p and 0xff0000) shr 16
                sir[1] = (p and 0x00ff00) shr 8
                sir[2] = (p and 0x0000ff)
                rbs = r1 - kotlin.math.abs(i2)
                rsum += sir[0] * rbs
                gsum += sir[1] * rbs
                bsum += sir[2] * rbs
                if (i2 > 0) {
                    rinsum += sir[0]
                    ginsum += sir[1]
                    binsum += sir[2]
                } else {
                    routsum += sir[0]
                    goutsum += sir[1]
                    boutsum += sir[2]
                }
            }
            stackpointer = radius

            for (x2 in 0 until w) {
                r[yi] = dv[rsum]
                g[yi] = dv[gsum]
                b[yi] = dv[bsum]

                rsum -= routsum
                gsum -= goutsum
                bsum -= boutsum

                stackstart = stackpointer - radius + div
                sir = stack[stackstart % div]

                routsum -= sir[0]
                goutsum -= sir[1]
                boutsum -= sir[2]

                if (y2 == 0) {
                    vmin[x2] = minOf(x2 + radius + 1, wm)
                }
                p = pix[yw + vmin[x2]]

                sir[0] = (p and 0xff0000) shr 16
                sir[1] = (p and 0x00ff00) shr 8
                sir[2] = (p and 0x0000ff)

                rinsum += sir[0]
                ginsum += sir[1]
                binsum += sir[2]

                rsum += rinsum
                gsum += ginsum
                bsum += binsum

                stackpointer = (stackpointer + 1) % div
                sir = stack[(stackpointer) % div]

                routsum += sir[0]
                goutsum += sir[1]
                boutsum += sir[2]

                rinsum -= sir[0]
                ginsum -= sir[1]
                binsum -= sir[2]

                yi++
            }
            yw += w
        }
        for (x2 in 0 until w) {
            rinsum = 0
            ginsum = 0
            binsum = 0
            routsum = 0
            goutsum = 0
            boutsum = 0
            rsum = 0
            gsum = 0
            bsum = 0
            yp = -radius * w
            for (i2 in -radius..radius) {
                yi = maxOf(0, yp) + x2

                sir = stack[i2 + radius]

                sir[0] = r[yi]
                sir[1] = g[yi]
                sir[2] = b[yi]

                rbs = r1 - kotlin.math.abs(i2)

                rsum += r[yi] * rbs
                gsum += g[yi] * rbs
                bsum += b[yi] * rbs

                if (i2 > 0) {
                    rinsum += sir[0]
                    ginsum += sir[1]
                    binsum += sir[2]
                } else {
                    routsum += sir[0]
                    goutsum += sir[1]
                    boutsum += sir[2]
                }

                if (i2 < hm) {
                    yp += w
                }
            }
            yi = x2
            stackpointer = radius
            for (y2 in 0 until h) {
                pix[yi] = (0xff000000.toInt() and pix[yi]) or (dv[rsum] shl 16) or (dv[gsum] shl 8) or dv[bsum]

                rsum -= routsum
                gsum -= goutsum
                bsum -= boutsum

                stackstart = stackpointer - radius + div
                sir = stack[stackstart % div]

                routsum -= sir[0]
                goutsum -= sir[1]
                boutsum -= sir[2]

                if (x2 == 0) {
                    vmin[y2] = minOf(y2 + r1, hm) * w
                }
                p = x2 + vmin[y2]

                sir[0] = r[p]
                sir[1] = g[p]
                sir[2] = b[p]

                rinsum += sir[0]
                ginsum += sir[1]
                binsum += sir[2]

                rsum += rinsum
                gsum += ginsum
                bsum += binsum

                stackpointer = (stackpointer + 1) % div
                sir = stack[stackpointer]

                routsum += sir[0]
                goutsum += sir[1]
                boutsum += sir[2]

                rinsum -= sir[0]
                ginsum -= sir[1]
                binsum -= sir[2]

                yi += w
            }
        }

        val result = Bitmap.createBitmap(w, h, Bitmap.Config.ARGB_8888)
        result.setPixels(pix, 0, w, 0, 0, w, h)
        return result
    }
}