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 android.net.Uri
import android.util.Log
import androidx.camera.core.ImageCapture
import androidx.camera.core.ImageCaptureException
import androidx.camera.core.ImageProxy
import androidx.exifinterface.media.ExifInterface
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.withContext
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
) {

    companion object {
        private const val TAG = "ImageCaptureHandler"
    }

    /**
     * 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 originalBitmap: Bitmap,
        val processedBitmap: 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) {
                        var currentBitmap: Bitmap? = null
                        try {
                            val imageRotation = imageProxy.imageInfo.rotationDegrees

                            currentBitmap = imageProxyToBitmap(imageProxy)
                            imageProxy.close()

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

                            currentBitmap = rotateBitmap(currentBitmap, imageRotation, isFrontCamera)

                            val processedBitmap = applyTiltShiftEffect(currentBitmap, blurParams)
                            // Keep originalBitmap alive — both are recycled after saving
                            val original = currentBitmap
                            currentBitmap = null

                            continuation.resume(ProcessedCapture(original, processedBitmap))
                        } catch (e: Exception) {
                            Log.e(TAG, "Image processing failed", e)
                            currentBitmap?.recycle()
                            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 both original and processed to disk (suspend-safe)
        if (captureResult is ProcessedCapture) {
            return try {
                val thumbnail = createThumbnail(captureResult.processedBitmap)
                val result = photoSaver.saveBitmapPair(
                    original = captureResult.originalBitmap,
                    processed = captureResult.processedBitmap,
                    orientation = ExifInterface.ORIENTATION_NORMAL,
                    location = location
                )
                if (result is SaveResult.Success) {
                    result.copy(thumbnail = thumbnail)
                } else {
                    thumbnail?.recycle()
                    result
                }
            } finally {
                captureResult.originalBitmap.recycle()
                captureResult.processedBitmap.recycle()
            }
        }

        return captureResult as SaveResult
    }

    /**
     * Creates a small thumbnail copy of a bitmap for in-app preview.
     */
    private fun createThumbnail(source: Bitmap, maxSize: Int = 160): Bitmap? {
        return try {
            val scale = maxSize.toFloat() / maxOf(source.width, source.height)
            val width = (source.width * scale).toInt()
            val height = (source.height * scale).toInt()
            Bitmap.createScaledBitmap(source, width, height, true)
        } catch (e: Exception) {
            Log.w(TAG, "Failed to create thumbnail", e)
            null
        }
    }

    /**
     * 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)
    }

    /**
     * Loads a gallery image and applies EXIF rotation, returning the bitmap for preview.
     * The caller owns the returned bitmap and is responsible for recycling it.
     */
    suspend fun loadGalleryImage(imageUri: Uri): Bitmap? = withContext(Dispatchers.IO) {
        try {
            val bitmap = loadBitmapFromUri(imageUri)
                ?: return@withContext null
            applyExifRotation(imageUri, bitmap)
        } catch (e: Exception) {
            Log.e(TAG, "Failed to load gallery image for preview", e)
            null
        }
    }

    /**
     * Processes an existing image from the gallery through the tilt-shift pipeline.
     * Loads the image, applies EXIF rotation, processes the effect, and saves both versions.
     */
    suspend fun processExistingImage(
        imageUri: Uri,
        blurParams: BlurParameters,
        location: Location?
    ): SaveResult = withContext(Dispatchers.IO) {
        var originalBitmap: Bitmap? = null
        var processedBitmap: Bitmap? = null
        try {
            originalBitmap = loadBitmapFromUri(imageUri)
                ?: return@withContext SaveResult.Error("Failed to load image")

            originalBitmap = applyExifRotation(imageUri, originalBitmap)

            processedBitmap = applyTiltShiftEffect(originalBitmap, blurParams)

            val thumbnail = createThumbnail(processedBitmap)

            val result = photoSaver.saveBitmapPair(
                original = originalBitmap,
                processed = processedBitmap,
                orientation = ExifInterface.ORIENTATION_NORMAL,
                location = location
            )

            if (result is SaveResult.Success) {
                result.copy(thumbnail = thumbnail)
            } else {
                thumbnail?.recycle()
                result
            }
        } catch (e: Exception) {
            Log.e(TAG, "Gallery image processing failed", e)
            SaveResult.Error("Processing failed: ${e.message}", e)
        } finally {
            originalBitmap?.recycle()
            processedBitmap?.recycle()
        }
    }

    /**
     * Loads a bitmap from a content URI.
     */
    private fun loadBitmapFromUri(uri: Uri): Bitmap? {
        return try {
            context.contentResolver.openInputStream(uri)?.use { stream ->
                BitmapFactory.decodeStream(stream)
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to load bitmap from URI", e)
            null
        }
    }

    /**
     * Reads EXIF orientation from a content URI and applies the
     * required rotation/flip to the bitmap.
     */
    private fun applyExifRotation(uri: Uri, bitmap: Bitmap): Bitmap {
        val orientation = try {
            context.contentResolver.openInputStream(uri)?.use { stream ->
                ExifInterface(stream).getAttributeInt(
                    ExifInterface.TAG_ORIENTATION,
                    ExifInterface.ORIENTATION_NORMAL
                )
            } ?: ExifInterface.ORIENTATION_NORMAL
        } catch (e: Exception) {
            Log.w(TAG, "Failed to read EXIF orientation", e)
            ExifInterface.ORIENTATION_NORMAL
        }

        val matrix = Matrix()
        when (orientation) {
            ExifInterface.ORIENTATION_ROTATE_90 -> matrix.postRotate(90f)
            ExifInterface.ORIENTATION_ROTATE_180 -> matrix.postRotate(180f)
            ExifInterface.ORIENTATION_ROTATE_270 -> matrix.postRotate(270f)
            ExifInterface.ORIENTATION_FLIP_HORIZONTAL -> matrix.postScale(-1f, 1f)
            ExifInterface.ORIENTATION_FLIP_VERTICAL -> matrix.postScale(1f, -1f)
            ExifInterface.ORIENTATION_TRANSPOSE -> {
                matrix.postRotate(90f)
                matrix.postScale(-1f, 1f)
            }
            ExifInterface.ORIENTATION_TRANSVERSE -> {
                matrix.postRotate(270f)
                matrix.postScale(-1f, 1f)
            }
            else -> return bitmap
        }

        val rotated = Bitmap.createBitmap(
            bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true
        )
        if (rotated != bitmap) {
            bitmap.recycle()
        }
        return rotated
    }

    /**
     * Applies tilt-shift blur effect to a bitmap.
     * Supports both linear and radial modes.
     *
     * All intermediate bitmaps are tracked and recycled in a finally block
     * so that an OOM or other exception does not leak native memory.
     */
    private fun applyTiltShiftEffect(source: Bitmap, params: BlurParameters): Bitmap {
        val width = source.width
        val height = source.height

        var result: Bitmap? = null
        var scaled: Bitmap? = null
        var blurred: Bitmap? = null
        var blurredFullSize: Bitmap? = null
        var mask: Bitmap? = null

        try {
            result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)

            val scaleFactor = 4
            val blurredWidth = width / scaleFactor
            val blurredHeight = height / scaleFactor

            scaled = Bitmap.createScaledBitmap(source, blurredWidth, blurredHeight, true)

            blurred = stackBlur(scaled, (params.blurAmount * 25).toInt().coerceIn(1, 25))
            scaled.recycle()
            scaled = null

            blurredFullSize = Bitmap.createScaledBitmap(blurred, width, height, true)
            blurred.recycle()
            blurred = null

            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()
            blurredFullSize = null
            mask.recycle()
            mask = null

            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)

            val output = result
            result = null // prevent finally from recycling the returned bitmap
            return output
        } finally {
            result?.recycle()
            scaled?.recycle()
            blurred?.recycle()
            blurredFullSize?.recycle()
            mask?.recycle()
        }
    }

    /**
     * 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
    }
}