/*
 * Hermes Node Android
 * Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
 * SPDX-License-Identifier: GPL-3.0-or-later
 */
package net.nexlab.hermesnodeandroid.capabilities

import android.content.Context
import android.graphics.Bitmap
import android.graphics.PixelFormat
import android.hardware.display.DisplayManager
import android.hardware.display.VirtualDisplay
import android.media.Image
import android.media.ImageReader
import android.media.projection.MediaProjectionManager
import android.os.Build
import android.os.Handler
import android.os.HandlerThread
import android.util.Base64
import android.util.DisplayMetrics
import android.view.WindowManager
import net.nexlab.hermesnodeandroid.AndroidRuntimeGrants
import net.nexlab.hermesnodeandroid.CapabilityResult
import org.json.JSONObject
import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer
import java.util.concurrent.ArrayBlockingQueue
import java.util.concurrent.TimeUnit
import kotlin.math.max
import kotlin.math.min

class MediaProjectionCapture(private val context: Context) {
    fun capture(params: JSONObject): CapabilityResult {
        val resultCode = AndroidRuntimeGrants.mediaProjectionResultCode
            ?: return CapabilityResult.unsupported("screenshot requires MediaProjection consent; tap 'Grant screenshot consent' in the Android app")
        val data = AndroidRuntimeGrants.mediaProjectionData
            ?: return CapabilityResult.unsupported("screenshot requires MediaProjection consent; tap 'Grant screenshot consent' in the Android app")

        val mgr = context.getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager
        val projection = mgr.getMediaProjection(resultCode, data)
            ?: return CapabilityResult.unsupported("MediaProjection grant is present but could not be activated; re-grant screenshot consent")

        val thread = HandlerThread("HermesScreenCapture").apply { start() }
        val handler = Handler(thread.looper)
        var virtualDisplay: VirtualDisplay? = null
        var imageReader: ImageReader? = null

        return try {
            val metrics = displayMetrics()
            val fullWidth = metrics.widthPixels
            val fullHeight = metrics.heightPixels
            val density = metrics.densityDpi
            imageReader = ImageReader.newInstance(fullWidth, fullHeight, PixelFormat.RGBA_8888, 2)
            val queue = ArrayBlockingQueue<Image>(1)
            imageReader.setOnImageAvailableListener({ reader ->
                reader.acquireLatestImage()?.let { queue.offer(it) }
            }, handler)

            virtualDisplay = projection.createVirtualDisplay(
                "HermesNodeCapture",
                fullWidth,
                fullHeight,
                density,
                DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
                imageReader.surface,
                null,
                handler,
            )

            val image = queue.poll(3, TimeUnit.SECONDS)
                ?: return CapabilityResult.unsupported("screenshot timed out waiting for a frame from MediaProjection")

            image.use {
                val cropped = cropBitmapFromImage(it, fullWidth, fullHeight, params)
                val out = ByteArrayOutputStream()
                val requestedFormat = params.optString("format", "png")
                val format = if (requestedFormat.equals("jpg", true) || requestedFormat.equals("jpeg", true)) {
                    Bitmap.CompressFormat.JPEG
                } else {
                    Bitmap.CompressFormat.PNG
                }
                val quality = min(100, max(1, params.optInt("quality", 90)))
                cropped.compress(format, quality, out)
                val bytes = out.toByteArray()
                return CapabilityResult.ok(
                    "format" to if (format == Bitmap.CompressFormat.JPEG) "jpg" else "png",
                    "encoding" to "base64",
                    "data" to Base64.encodeToString(bytes, Base64.NO_WRAP),
                    "bytes" to bytes.size,
                    "width" to cropped.width,
                    "height" to cropped.height,
                    "full_width" to fullWidth,
                    "full_height" to fullHeight,
                )
            }
        } catch (e: SecurityException) {
            CapabilityResult.unsupported("screenshot permission denied: ${e.message}")
        } catch (e: Exception) {
            CapabilityResult.unsupported("screenshot failed: ${e.message}")
        } finally {
            runCatching { imageReader?.setOnImageAvailableListener(null, null) }
            runCatching { virtualDisplay?.release() }
            runCatching { projection.stop() }
            thread.quitSafely()
        }
    }

    private fun cropBitmapFromImage(image: Image, fullWidth: Int, fullHeight: Int, params: JSONObject): Bitmap {
        val plane = image.planes[0]
        val buffer: ByteBuffer = plane.buffer
        val pixelStride = plane.pixelStride
        val rowStride = plane.rowStride
        val rowPadding = rowStride - pixelStride * fullWidth
        val bitmap = Bitmap.createBitmap(fullWidth + rowPadding / pixelStride, fullHeight, Bitmap.Config.ARGB_8888)
        bitmap.copyPixelsFromBuffer(buffer)
        val normalized = Bitmap.createBitmap(bitmap, 0, 0, fullWidth, fullHeight)
        if (bitmap !== normalized) bitmap.recycle()

        val requestedX = params.optInt("x", 0)
        val requestedY = params.optInt("y", 0)
        val requestedWidth = params.optInt("width", fullWidth)
        val requestedHeight = params.optInt("height", fullHeight)
        val x = requestedX.coerceIn(0, max(0, fullWidth - 1))
        val y = requestedY.coerceIn(0, max(0, fullHeight - 1))
        val width = requestedWidth.coerceIn(1, fullWidth - x)
        val height = requestedHeight.coerceIn(1, fullHeight - y)
        if (x == 0 && y == 0 && width == fullWidth && height == fullHeight) return normalized
        val cropped = Bitmap.createBitmap(normalized, x, y, width, height)
        normalized.recycle()
        return cropped
    }

    private fun displayMetrics(): DisplayMetrics {
        val metrics = DisplayMetrics()
        @Suppress("DEPRECATION")
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
            val display = context.display ?: context.getSystemService(WindowManager::class.java)?.defaultDisplay
            display?.getRealMetrics(metrics)
        } else {
            val wm = context.getSystemService(Context.WINDOW_SERVICE) as WindowManager
            wm.defaultDisplay.getRealMetrics(metrics)
        }
        if (metrics.widthPixels <= 0 || metrics.heightPixels <= 0) {
            metrics.setTo(context.resources.displayMetrics)
        }
        return metrics
    }
}
