package com.ledgrab.android

import android.graphics.Bitmap
import android.graphics.PixelFormat
import android.hardware.display.DisplayManager
import android.hardware.display.VirtualDisplay
import android.media.ImageReader
import android.media.projection.MediaProjection
import android.os.Handler
import android.os.HandlerThread
import android.util.DisplayMetrics
import android.util.Log
import java.nio.ByteBuffer

/**
 * Captures the Android screen via MediaProjection and feeds frames
 * to [PythonBridge].
 *
 * Frames are downscaled to [targetWidth] x [targetHeight] before
 * crossing the JNI boundary to minimize overhead.  For LED ambient
 * lighting, even 480x270 contains far more data than needed.
 */
class ScreenCapture(
    private val projection: MediaProjection,
    private val metrics: DisplayMetrics,
    private val bridge: PythonBridge,
    private val targetWidth: Int = 480,
    private val targetHeight: Int = 270,
    private val targetFps: Int = 30,
    private val onProjectionStopped: () -> Unit = {},
) {
    companion object {
        private const val TAG = "ScreenCapture"
        private const val VIRTUAL_DISPLAY_NAME = "LedGrabCapture"
    }

    private var virtualDisplay: VirtualDisplay? = null
    private var imageReader: ImageReader? = null
    private var captureThread: HandlerThread? = null
    private var captureHandler: Handler? = null
    @Volatile private var running = false
    private var lastFrameTimeMs = 0L
    private val frameIntervalMs = 1000L / targetFps

    /**
     * Start capturing the screen.
     */
    fun start() {
        if (running) return
        running = true

        captureThread = HandlerThread("LedGrab-Capture").also { it.start() }
        captureHandler = Handler(captureThread!!.looper)

        // Android 14+ requires registering a callback before createVirtualDisplay
        projection.registerCallback(object : MediaProjection.Callback() {
            override fun onStop() {
                Log.i(TAG, "MediaProjection stopped (external)")
                stop()
                // Notify the service so the foreground notification /
                // Python server get torn down too — otherwise a stale
                // "Running" notification lingers after the user taps
                // Android's system Cast/Screen-capture stop banner.
                onProjectionStopped()
            }
        }, captureHandler)

        imageReader = ImageReader.newInstance(
            targetWidth,
            targetHeight,
            PixelFormat.RGBA_8888,
            2,  // maxImages — double buffer
        )

        imageReader?.setOnImageAvailableListener({ reader ->
            if (!running) return@setOnImageAvailableListener

            val now = System.currentTimeMillis()
            if (now - lastFrameTimeMs < frameIntervalMs) {
                // Skip frame to maintain target FPS
                reader.acquireLatestImage()?.close()
                return@setOnImageAvailableListener
            }

            val image = reader.acquireLatestImage() ?: return@setOnImageAvailableListener
            try {
                val plane = image.planes[0]
                val buffer = plane.buffer
                val rowStride = plane.rowStride
                val pixelStride = plane.pixelStride

                // Handle row padding: rowStride may be > width * pixelStride
                val rgbaBytes = if (rowStride == targetWidth * pixelStride) {
                    // No padding — direct copy
                    val bytes = ByteArray(buffer.remaining())
                    buffer.get(bytes)
                    bytes
                } else {
                    // Strip row padding
                    val rowBytes = targetWidth * pixelStride
                    val bytes = ByteArray(targetWidth * targetHeight * 4)
                    for (row in 0 until targetHeight) {
                        buffer.position(row * rowStride)
                        buffer.get(bytes, row * rowBytes, rowBytes)
                    }
                    bytes
                }

                bridge.pushFrame(rgbaBytes, targetWidth, targetHeight)
                lastFrameTimeMs = now
            } catch (e: Exception) {
                Log.w(TAG, "Frame processing error: ${e.message}")
            } finally {
                image.close()
            }
        }, captureHandler)

        virtualDisplay = projection.createVirtualDisplay(
            VIRTUAL_DISPLAY_NAME,
            targetWidth,
            targetHeight,
            metrics.densityDpi,
            DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
            imageReader?.surface,
            null,
            captureHandler,
        )

        Log.i(TAG, "Screen capture started (${targetWidth}x${targetHeight} @ ${targetFps}fps)")
    }

    /**
     * Stop capturing and release all resources.
     */
    fun stop() {
        running = false
        // Order matters: detach the listener BEFORE releasing the
        // VirtualDisplay so the handler can't be re-entered with stale
        // resources, then quit & join the handler thread, only then
        // close the ImageReader.
        runCatching { imageReader?.setOnImageAvailableListener(null, null) }
        runCatching { virtualDisplay?.release() }
        virtualDisplay = null

        captureThread?.quitSafely()
        runCatching { captureThread?.join(500) }
        captureThread = null
        captureHandler = null

        runCatching { imageReader?.close() }
        imageReader = null

        Log.i(TAG, "Screen capture stopped")
    }
}
