package com.ledgrab.android

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.os.SystemClock
import android.util.DisplayMetrics
import android.util.Log

/**
 * Captures the Android screen via MediaProjection and feeds frames
 * to [PythonBridge].
 *
 * Frames are downscaled to roughly [targetWidth] x [targetHeight] before
 * crossing the JNI boundary to minimize overhead. The actual capture
 * dimensions preserve the source screen's aspect ratio (snapped to even
 * pixels for codec friendliness) so non-16:9 displays don't get
 * squashed.
 */
class ScreenCapture(
    private val projection: MediaProjection,
    private val metrics: DisplayMetrics,
    private val bridge: PythonBridge,
    targetWidth: Int = 480,
    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"
    }

    // Snap to the source aspect ratio so we don't squash 21:9 / portrait
    // / rotated screens. Width is the budget; height follows.
    private val captureWidth: Int
    private val captureHeight: Int

    init {
        val srcW = metrics.widthPixels.coerceAtLeast(1).toFloat()
        val srcH = metrics.heightPixels.coerceAtLeast(1).toFloat()
        val budget = targetWidth.coerceAtLeast(16)
        val aspect = srcW / srcH
        val w = budget
        val h = (w / aspect).toInt().coerceAtLeast(16)
        // Bias toward even dimensions — some encoders/ImageReaders are
        // unhappy with odd sizes when row strides come into play.
        captureWidth = (w and 1.inv()).coerceAtLeast(16)
        captureHeight = (h and 1.inv()).coerceAtLeast(16)
        if (captureWidth != targetWidth || captureHeight != targetHeight) {
            Log.i(
                TAG,
                "Capture size adjusted for ${srcW.toInt()}x${srcH.toInt()} " +
                    "(${"%.2f".format(aspect)}:1) → ${captureWidth}x$captureHeight",
            )
        }
    }

    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

    // Reusable RGBA frame buffer — sized once for the capture dimensions.
    // The capture handler is single-threaded so no synchronisation is
    // required around this buffer (each callback runs to completion
    // before the next is dispatched). Eliminates ~15 MB/s of per-frame
    // garbage at 30 fps × 480×270×4 B that previously caused GC pauses
    // on low-end TV boxes.
    private val frameBuffer: ByteArray = ByteArray(captureWidth * captureHeight * 4)

    // Monotonic frame pacing. `nextFrameNanos` is the target render
    // time of the next frame; carrying it forward as an accumulator
    // avoids the integer-division drift the wall-clock version had
    // (e.g. 30 fps → 33 ms produced ~30.3 fps).
    private val frameIntervalNanos = (1_000_000_000L / targetFps.coerceAtLeast(1))
    private var nextFrameNanos = 0L

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

        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)")
                // We're on captureHandler's thread here — calling stop()
                // directly would self-join captureThread (handler.join()
                // from inside the handler thread hangs until the join
                // timeout, then closes resources while we're STILL
                // inside this callback). Just flip `running` to halt
                // frame processing and hand off to the service; its
                // onDestroy will call stop() from the main thread,
                // which is safe to join captureThread from.
                running = false
                onProjectionStopped()
            }
        }, captureHandler)

        imageReader = ImageReader.newInstance(
            captureWidth,
            captureHeight,
            PixelFormat.RGBA_8888,
            3,  // maxImages — small ring buffer; 3 is more forgiving than 2 under jitter
        )

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

            val now = SystemClock.elapsedRealtimeNanos()
            if (now < nextFrameNanos) {
                // Too early — drop this image to stay on cadence.
                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
                val rowBytes = captureWidth * pixelStride
                val expected = rowBytes * captureHeight

                // Fill the reusable buffer. Two paths:
                //   - rowStride == rowBytes: bulk get into the buffer
                //   - rowStride >  rowBytes: row-by-row copy stripping padding
                if (rowStride == rowBytes && buffer.remaining() >= expected) {
                    buffer.get(frameBuffer, 0, expected)
                } else {
                    for (row in 0 until captureHeight) {
                        buffer.position(row * rowStride)
                        buffer.get(frameBuffer, row * rowBytes, rowBytes)
                    }
                }

                bridge.pushFrame(frameBuffer, captureWidth, captureHeight)

                // Advance the pacing accumulator. If we fell badly behind
                // (long GC, JNI stall), snap forward to "now" instead of
                // accumulating a burst of catch-up frames.
                nextFrameNanos += frameIntervalNanos
                if (now - nextFrameNanos > frameIntervalNanos * 4) {
                    nextFrameNanos = now + frameIntervalNanos
                }
            } catch (e: Exception) {
                Log.w(TAG, "Frame processing error: ${e.message}")
            } finally {
                image.close()
            }
        }, captureHandler)

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

        Log.i(TAG, "Screen capture started (${captureWidth}x${captureHeight} @ ${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")
    }
}
