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 no.naiv.tiltshift.util.StackBlur
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"
        /** Maximum decoded image dimension to prevent OOM from huge gallery images. */
        private const val MAX_IMAGE_DIMENSION = 4096
        /** Scale factor for downscaling blur and mask computations. */
        private const val SCALE_FACTOR = 4
    }

    /**
     * Type-safe outcome of the camera capture callback.
     * Eliminates unsafe `as Any` / `as SaveResult` casts.
     */
    private sealed class CaptureOutcome {
        class Processed(val original: Bitmap, val processed: Bitmap) : CaptureOutcome()
        class Failed(val result: SaveResult.Error) : CaptureOutcome()
    }

    /**
     * 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<CaptureOutcome> { 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)

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

                            currentBitmap = rotateBitmap(currentBitmap, imageRotation, isFrontCamera)

                            val processedBitmap = applyTiltShiftEffect(currentBitmap, blurParams)
                            val original = currentBitmap
                            currentBitmap = null

                            continuation.resume(CaptureOutcome.Processed(original, processedBitmap))
                        } catch (e: Exception) {
                            Log.e(TAG, "Image processing failed", e)
                            currentBitmap?.recycle()
                            continuation.resume(
                                CaptureOutcome.Failed(SaveResult.Error("Failed to process image. Please try again.", e))
                            )
                        } finally {
                            imageProxy.close()
                        }
                    }

                    override fun onError(exception: ImageCaptureException) {
                        Log.e(TAG, "Image capture failed", exception)
                        continuation.resume(
                            CaptureOutcome.Failed(SaveResult.Error("Failed to capture photo. Please try again.", exception))
                        )
                    }
                }
            )
        }

        // Phase 2: save both original and processed to disk (suspend-safe)
        return when (captureResult) {
            is CaptureOutcome.Failed -> captureResult.result
            is CaptureOutcome.Processed -> {
                var thumbnail: Bitmap? = null
                try {
                    thumbnail = createThumbnail(captureResult.processed)
                    val result = photoSaver.saveBitmapPair(
                        original = captureResult.original,
                        processed = captureResult.processed,
                        orientation = ExifInterface.ORIENTATION_NORMAL,
                        location = location
                    )
                    if (result is SaveResult.Success) {
                        val output = result.copy(thumbnail = thumbnail)
                        thumbnail = null // prevent finally from recycling the returned thumbnail
                        output
                    } else {
                        result
                    }
                } finally {
                    thumbnail?.recycle()
                    captureResult.original.recycle()
                    captureResult.processed.recycle()
                }
            }
        }
    }

    /**
     * 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 the result.
     */
    suspend fun processExistingImage(
        imageUri: Uri,
        blurParams: BlurParameters,
        location: Location?
    ): SaveResult = withContext(Dispatchers.IO) {
        var originalBitmap: Bitmap? = null
        var processedBitmap: Bitmap? = null
        var thumbnail: Bitmap? = null
        try {
            originalBitmap = loadBitmapFromUri(imageUri)
                ?: return@withContext SaveResult.Error("Failed to load image")

            originalBitmap = applyExifRotation(imageUri, originalBitmap)

            processedBitmap = applyTiltShiftEffect(originalBitmap, blurParams)

            thumbnail = createThumbnail(processedBitmap)

            val result = photoSaver.saveBitmap(
                bitmap = processedBitmap,
                orientation = ExifInterface.ORIENTATION_NORMAL,
                location = location
            )

            if (result is SaveResult.Success) {
                val output = result.copy(thumbnail = thumbnail)
                thumbnail = null // prevent finally from recycling the returned thumbnail
                output
            } else {
                result
            }
        } catch (e: SecurityException) {
            Log.e(TAG, "Permission denied while processing gallery image", e)
            SaveResult.Error("Permission denied. Please grant access and try again.", e)
        } catch (e: Exception) {
            Log.e(TAG, "Gallery image processing failed", e)
            SaveResult.Error("Failed to process image. Please try again.", e)
        } finally {
            thumbnail?.recycle()
            originalBitmap?.recycle()
            processedBitmap?.recycle()
        }
    }

    /**
     * Loads a bitmap from a content URI with dimension bounds checking
     * to prevent OOM from extremely large images.
     */
    private fun loadBitmapFromUri(uri: Uri): Bitmap? {
        return try {
            // First pass: read dimensions without decoding pixels
            val options = BitmapFactory.Options().apply { inJustDecodeBounds = true }
            context.contentResolver.openInputStream(uri)?.use { stream ->
                BitmapFactory.decodeStream(stream, null, options)
            } ?: run {
                Log.e(TAG, "Could not open input stream for URI (dimensions pass): $uri")
                return null
            }

            if (options.outWidth <= 0 || options.outHeight <= 0) {
                Log.e(TAG, "Image has invalid dimensions: ${options.outWidth}x${options.outHeight}, mime: ${options.outMimeType}")
                return null
            }

            // Calculate sample size to stay within MAX_IMAGE_DIMENSION
            val maxDim = maxOf(options.outWidth, options.outHeight)
            val sampleSize = if (maxDim > MAX_IMAGE_DIMENSION) {
                var sample = 1
                while (maxDim / sample > MAX_IMAGE_DIMENSION) sample *= 2
                sample
            } else 1

            // Second pass: decode with sample size
            val decodeOptions = BitmapFactory.Options().apply { inSampleSize = sampleSize }
            val bitmap = context.contentResolver.openInputStream(uri)?.use { stream ->
                BitmapFactory.decodeStream(stream, null, decodeOptions)
            }

            if (bitmap == null) {
                Log.e(TAG, "BitmapFactory.decodeStream returned null for URI: $uri (mime: ${options.outMimeType})")
            }

            bitmap
        } catch (e: SecurityException) {
            Log.e(TAG, "Permission denied loading bitmap from URI", e)
            null
        } 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 effect to a bitmap for real-time preview.
     * Runs on [Dispatchers.IO]. The caller owns the returned bitmap.
     */
    suspend fun applyTiltShiftPreview(source: Bitmap, params: BlurParameters): Bitmap =
        withContext(Dispatchers.IO) {
            applyTiltShiftEffect(source, params)
        }

    /**
     * Applies tilt-shift blur effect to a bitmap.
     * Supports both linear and radial modes.
     *
     * The gradient mask is computed at 1/4 resolution (matching the blur downscale)
     * and upscaled for compositing, reducing peak memory by ~93%.
     *
     * 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

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

            val blurredWidth = width / SCALE_FACTOR
            val blurredHeight = height / SCALE_FACTOR

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

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

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

            // Compute mask at reduced resolution and upscale to avoid full-res per-pixel trig
            val maskPixels = createGradientMaskPixels(blurredWidth, blurredHeight, params)
            val fullMaskPixels = upscaleMask(maskPixels, blurredWidth, blurredHeight, width, height)

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

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

            blurredFullSize.recycle()
            blurredFullSize = null

            for (i in pixels.indices) {
                val maskAlpha = (fullMaskPixels[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()
        }
    }

    /**
     * Creates a gradient mask as a pixel array at the given dimensions.
     * Returns packed ARGB ints where the blue channel encodes blur amount.
     */
    private fun createGradientMaskPixels(width: Int, height: Int, params: BlurParameters): IntArray {
        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 -> {
                        val dx = x - centerX
                        val dy = y - centerY
                        val rotatedY = -dx * sinAngle + dy * cosAngle
                        kotlin.math.abs(rotatedY)
                    }
                    BlurMode.RADIAL -> {
                        var dx = x - centerX
                        var dy = y - centerY
                        dx *= screenAspect
                        val rotatedX = dx * cosAngle - dy * sinAngle
                        val rotatedY = dx * sinAngle + dy * cosAngle
                        val adjustedX = rotatedX / params.aspectRatio
                        sqrt(adjustedX * adjustedX + rotatedY * rotatedY)
                    }
                }

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

        return pixels
    }

    /**
     * Bilinear upscale of a mask pixel array from small dimensions to full dimensions.
     */
    private fun upscaleMask(
        smallPixels: IntArray,
        smallW: Int, smallH: Int,
        fullW: Int, fullH: Int
    ): IntArray {
        val fullPixels = IntArray(fullW * fullH)
        val xRatio = smallW.toFloat() / fullW
        val yRatio = smallH.toFloat() / fullH

        for (y in 0 until fullH) {
            val srcY = y * yRatio
            val y0 = srcY.toInt().coerceIn(0, smallH - 1)
            val y1 = (y0 + 1).coerceIn(0, smallH - 1)
            val yFrac = srcY - y0

            for (x in 0 until fullW) {
                val srcX = x * xRatio
                val x0 = srcX.toInt().coerceIn(0, smallW - 1)
                val x1 = (x0 + 1).coerceIn(0, smallW - 1)
                val xFrac = srcX - x0

                // Bilinear interpolation on the blue channel (all channels are equal)
                val v00 = smallPixels[y0 * smallW + x0] and 0xFF
                val v10 = smallPixels[y0 * smallW + x1] and 0xFF
                val v01 = smallPixels[y1 * smallW + x0] and 0xFF
                val v11 = smallPixels[y1 * smallW + x1] and 0xFF

                val top = v00 + (v10 - v00) * xFrac
                val bottom = v01 + (v11 - v01) * xFrac
                val gray = (top + (bottom - top) * yFrac).toInt().coerceIn(0, 255)

                fullPixels[y * fullW + x] = (0xFF shl 24) or (gray shl 16) or (gray shl 8) or gray
            }
        }

        return fullPixels
    }

}