package deckers.thibault.aves.decoding

import android.content.Context
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.BitmapRegionDecoder
import android.graphics.ColorSpace
import android.graphics.Rect
import android.net.Uri
import android.os.Build
import android.util.Log
import androidx.core.graphics.createBitmap
import com.bumptech.glide.Glide
import deckers.thibault.aves.channel.streams.darttoplatform.ByteSink
import deckers.thibault.aves.glide.AvesAppGlideModule
import deckers.thibault.aves.glide.MultiPageImage
import deckers.thibault.aves.utils.BitmapRegionDecoderCompat
import deckers.thibault.aves.utils.BitmapUtils
import deckers.thibault.aves.utils.BitmapUtils.describe
import deckers.thibault.aves.utils.LogUtils
import deckers.thibault.aves.utils.MathUtils
import deckers.thibault.aves.utils.MemoryUtils
import deckers.thibault.aves.utils.MimeTypes
import deckers.thibault.aves.utils.StorageUtils
import java.io.ByteArrayInputStream
import java.nio.ByteBuffer
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock
import kotlin.math.max
import kotlin.math.roundToInt

// As of Android 14 (API 34), `BitmapRegionDecoder` documentation states
// that "only the JPEG, PNG, WebP and HEIF formats are supported"
// but in practice it successfully decodes some others.
class RegionFetcher internal constructor(
    private val context: Context,
) {
    suspend fun fetch(
        uri: Uri,
        pageId: Int?,
        decoded: Boolean,
        mimeType: String,
        sampleSize: Int,
        regionRect: Rect,
        imageWidth: Int,
        imageHeight: Int,
        requestKey: Pair<Uri, Int?> = Pair(uri, pageId),
        result: ByteSink,
    ) {
        if (pageId != null && MultiPageImage.isSupported(mimeType)) {
            // use export for requested page
            val exportUri = exportUris.getOrPut(requestKey) { createTemporaryExport(uri, mimeType, pageId) }
            fetch(
                uri = exportUri,
                pageId = null,
                decoded = decoded,
                mimeType = EXPORT_MIME_TYPE,
                sampleSize = sampleSize,
                regionRect = regionRect,
                imageWidth = imageWidth,
                imageHeight = imageHeight,
                requestKey = requestKey,
                result = result,
            )
            return
        }

        try {
            val decoder = getOrCreateDecoder(context, uri, requestKey)
            if (decoder == null) {
                result.error("fetch-read-null", "failed to open file for mimeType=$mimeType uri=$uri regionRect=$regionRect", null)
                return
            }

            // with raw images, the known image size may not match the decoded image size
            // so we scale the requested region accordingly
            var effectiveRect = regionRect
            var effectiveSampleSize = sampleSize

            if (imageWidth != decoder.width || imageHeight != decoder.height) {
                val xf = decoder.width.toDouble() / imageWidth
                val yf = decoder.height.toDouble() / imageHeight
                effectiveRect = Rect(
                    (regionRect.left * xf).roundToInt(),
                    (regionRect.top * yf).roundToInt(),
                    (regionRect.right * xf).roundToInt(),
                    (regionRect.bottom * yf).roundToInt(),
                )
                val factor = MathUtils.highestPowerOf2((1 / max(xf, yf)).roundToInt())
                if (factor > 1) {
                    effectiveSampleSize = max(1, effectiveSampleSize / factor)
                }
            }

            val options = BitmapFactory.Options().apply {
                inSampleSize = effectiveSampleSize
                // Specifying preferred config and color space avoids the need for conversion afterwards,
                // but may prevent decoding (e.g. from RGBA_1010102 to ARGB_8888 on some devices).
                inPreferredConfig = PREFERRED_CONFIG
                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
                    inPreferredColorSpace = ColorSpace.get(ColorSpace.Named.SRGB)
                }
            }

            val pixelCount = effectiveRect.width() * effectiveRect.height() / effectiveSampleSize
            val targetBitmapSizeBytes = BitmapUtils.getExpectedImageSize(pixelCount.toLong(), options.inPreferredConfig)
            if (!MemoryUtils.canAllocate(targetBitmapSizeBytes)) {
                // decoding a region that large would yield an OOM when creating the bitmap
                result.error("fetch-large-region", "Region too large for uri=$uri regionRect=$regionRect", null)
                return
            }

            var bitmap = decoder.decodeRegion(effectiveRect, options)
            if (bitmap == null) {
                // retry without specifying config or color space,
                // falling back to custom byte conversion afterwards
                options.inPreferredConfig = null
                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O && options.inPreferredColorSpace != null) {
                    options.inPreferredColorSpace = null
                }
                bitmap = decoder.decodeRegion(effectiveRect, options)
            }

            val bytes = BitmapUtils.getBytes(bitmap, recycle = true, decoded = decoded, mimeType)
            if (bytes == null) {
                result.error("fetch-null", "failed to decode region for uri=$uri regionRect=$regionRect", null)
            } else {
                result.streamBytes(ByteArrayInputStream(bytes))
            }
        } catch (e: Exception) {
            if (EXPORT_MIME_TYPE != mimeType) {
                // retry with export on failure,
                // as some formats are not fully supported by `BitmapRegionDecoder`
                val exportUri = exportUris.getOrPut(requestKey) { createTemporaryExport(uri, mimeType, pageId) }
                fetch(
                    uri = exportUri,
                    pageId = null,
                    decoded = decoded,
                    mimeType = EXPORT_MIME_TYPE,
                    sampleSize = sampleSize,
                    regionRect = regionRect,
                    imageWidth = imageWidth,
                    imageHeight = imageHeight,
                    requestKey = requestKey,
                    result = result,
                )
                return
            }

            result.error("fetch-read-exception", "failed to initialize region decoder for uri=$uri regionRect=$regionRect", e.message)
        }
    }

    private suspend fun createTemporaryExport(uri: Uri, mimeType: String, pageId: Int?): Uri {
        val exportFormat = EXPORT_FORMAT
        Log.d(LOG_TAG, "create export for uri=$uri mimeType=$mimeType pageId=$pageId exportFormat=$exportFormat")
        val target = Glide.with(context)
            .asBitmap()
            .apply(AvesAppGlideModule.uncachedFullImageOptions)
            .load(AvesAppGlideModule.getModel(context, uri, mimeType, pageId))
            .submit()

        try {
            val bitmap = target.get()
            val tempFile = StorageUtils.createTempFile(context).apply {
                outputStream().use { output ->
                    val encodedExport = bitmap.compress(exportFormat, 100, output)
                    if (!encodedExport) {
                        Log.w(LOG_TAG, "failed export via encoded bytes for uri=$uri mimeType=$mimeType pageId=$pageId exportFormat=$exportFormat, with bitmap=${bitmap.describe()}")

                        val decodedBytes = BitmapUtils.getBytes(bitmap, recycle = false, decoded = true, mimeType)
                        if (decodedBytes != null) {
                            val exportBitmap = createBitmap(bitmap.width, bitmap.height, PREFERRED_CONFIG)
                            exportBitmap.copyPixelsFromBuffer(ByteBuffer.wrap(decodedBytes))
                            val decodedExport = exportBitmap.compress(exportFormat, 100, output)
                            if (!decodedExport) {
                                Log.w(LOG_TAG, "failed to compress exportBitmap=${bitmap.describe()}")
                            }
                        }
                    }
                }
            }
            return Uri.fromFile(tempFile)
        } finally {
            Glide.with(context).clear(target)
        }
    }

    private data class DecoderRef(
        val requestKey: Pair<Uri, Int?>,
        val decoder: BitmapRegionDecoder,
    )

    companion object {
        private val LOG_TAG = LogUtils.createTag<RegionFetcher>()
        private val PREFERRED_CONFIG = Bitmap.Config.ARGB_8888
        private const val DECODER_POOL_SIZE = 3
        private const val EXPORT_MIME_TYPE = MimeTypes.JPEG
        private val EXPORT_FORMAT = Bitmap.CompressFormat.JPEG
        private val decoderPool = ArrayList<DecoderRef>()
        private val exportUris = HashMap<Pair<Uri, Int?>, Uri>()

        private val poolLock = ReentrantLock()

        private fun getOrCreateDecoder(context: Context, uri: Uri, requestKey: Pair<Uri, Int?>): BitmapRegionDecoder? {
            poolLock.withLock {
                var decoderRef = decoderPool.firstOrNull { it.requestKey == requestKey }
                if (decoderRef == null) {
                    val newDecoder = StorageUtils.openInputStream(context, uri)?.use { input ->
                        BitmapRegionDecoderCompat.newInstance(input)
                    }
                    if (newDecoder == null) {
                        return null
                    }
                    decoderRef = DecoderRef(requestKey, newDecoder)
                } else {
                    decoderPool.remove(decoderRef)
                }
                decoderPool.add(0, decoderRef)
                while (decoderPool.size > DECODER_POOL_SIZE) {
                    decoderPool.removeAt(decoderPool.size - 1)
                }
                return decoderRef.decoder
            }
        }
    }
}
