package com.ledgrab.android

import android.graphics.PixelFormat
import android.media.ImageReader
import android.media.MediaCodec
import android.media.MediaFormat
import android.os.Handler
import android.os.HandlerThread
import android.util.Log
import java.io.InputStream
import java.util.concurrent.atomic.AtomicInteger

/**
 * Root-only screen capture backed by `/system/bin/screenrecord`.
 *
 * When the device is rooted (Magisk), we spawn ``su -c screenrecord
 * --output-format=h264 …`` and read the encoder's H.264 bitstream from
 * stdout. A MediaCodec decoder is configured to render into an
 * ImageReader's Surface, and the ImageReader's `onImageAvailable`
 * callback feeds RGBA bytes back to the Python pipeline — same sink as
 * the MediaProjection path, just a different source.
 *
 * Why bother: MediaProjection forces Android to draw a persistent
 * capture indicator overlay (unavoidable on stock Android 14+). The
 * root path sidesteps that entirely because `screenrecord` runs as
 * system UID through `su`. It also skips the one-time MediaProjection
 * consent dialog, which matters on TV-only setups where keyboard input
 * is awkward.
 */
class RootScreenrecord(
    private val bridge: PythonBridge,
    private val width: Int = 480,
    private val height: Int = 270,
    private val bitRate: Int = 4_000_000,
    private val fps: Int = 30,
) {
    companion object {
        private const val TAG = "RootScreenrecord"
        private const val MIME_TYPE = MediaFormat.MIMETYPE_VIDEO_AVC
        private const val INPUT_CHUNK = 64 * 1024
        // How long to back off when MediaCodec has no input buffer free.
        // 50 ms keeps the input pump from busy-spinning if the decoder
        // is stalled (codec init, severe stall, etc.).
        private const val NO_BUFFER_BACKOFF_MS = 5L
    }

    // Instance is single-use: stop() permanently disposes it. Callers
    // wanting to restart the pipeline must construct a new instance —
    // see CaptureService.restartRootPipeline().
    @Volatile private var process: Process? = null
    private var decoder: MediaCodec? = null
    private var imageReader: ImageReader? = null
    private var readerThread: HandlerThread? = null
    private var inputThread: Thread? = null
    private var outputThread: Thread? = null
    @Volatile private var running = false
    private val framesDeliveredCounter = AtomicInteger(0)
    // disposed gates duplicate-stop calls only — not start() after
    // stop() (which is unsupported, see note above). Set at the START
    // of cleanup so a second concurrent stop() (rare under @Synchronized
    // but possible if a future caller drops it) doesn't re-run runCatching
    // blocks against already-released resources.
    @Volatile private var disposed = false
    // Guards process respawn vs. concurrent disposal. The input pump
    // can spawn a fresh `su -c screenrecord` after EOF; without this
    // lock, stop() could destroy the OLD process between spawn and
    // assignment, leaving the new one orphaned (GPU encoder leak).
    private val processLock = Any()

    // Reusable RGBA buffer for ImageReader callbacks (single-threaded
    // reader callback). See ScreenCapture for the rationale: avoids
    // ~15 MB/s of per-frame garbage at 30 fps × 480×270×4 B.
    private val frameBuffer: ByteArray = ByteArray(width * height * 4)

    /** Monotonic count of frames pushed to the Python bridge. */
    val framesDelivered: Int get() = framesDeliveredCounter.get()

    /** True once at least one frame has reached the Python bridge. */
    val hasProducedFrame: Boolean get() = framesDelivered > 0

    /**
     * Start the capture pipeline. Returns false if `su`/`screenrecord`
     * couldn't be launched at all; true doesn't guarantee frames will
     * actually flow (the caller should monitor [hasProducedFrame] and
     * be ready to fall back to MediaProjection on timeout).
     */
    fun start(): Boolean {
        if (running) return true
        running = true

        try {
            imageReader = buildImageReader()
            decoder = buildDecoder(imageReader!!)
            process = spawnScreenrecord() ?: run {
                stop()
                return false
            }
            startInputPump(process!!.inputStream, decoder!!)
            startOutputDrain(decoder!!)
            Log.i(TAG, "Root capture pipeline started (${width}x$height @ ${fps}fps)")
            return true
        } catch (e: Exception) {
            Log.e(TAG, "Failed to start root capture", e)
            stop()
            return false
        }
    }

    /** Stop everything and release resources. Idempotent. Single-use: do not call start() again. */
    @Synchronized
    fun stop() {
        if (disposed) return
        disposed = true
        // Order matters: signal first so worker loops drop out, then
        // stop the codec on the thread that created it (this one), then
        // join workers BEFORE releasing the codec/ImageReader they may
        // still be touching, then kill the external screenrecord process.
        running = false

        runCatching { decoder?.stop() }

        inputThread?.interrupt()
        outputThread?.interrupt()
        runCatching { inputThread?.join(500) }
        runCatching { outputThread?.join(500) }
        inputThread = null
        outputThread = null

        // Best-effort: kill the screenrecord child before reaping `su`,
        // otherwise screenrecord can outlive su as an orphan and keep
        // the GPU encoder busy. Fire-and-forget; ignore failures.
        runCatching {
            Root.runAsRoot(arrayOf("pkill", "-TERM", "screenrecord"), timeoutSeconds = 2)
        }

        runCatching { decoder?.release() }
        decoder = null

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

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

        // Use the same lock as the respawn path so we don't destroy a
        // not-yet-published process or leak one that was spawned after
        // we already destroyed the old reference.
        synchronized(processLock) {
            runCatching { process?.destroy() }
            process = null
        }

        Log.i(TAG, "Root capture pipeline stopped (frames delivered: ${framesDelivered})")
    }

    private fun buildImageReader(): ImageReader {
        val thread = HandlerThread("RootReader").apply { start() }
        readerThread = thread
        val handler = Handler(thread.looper)

        val reader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888, 3)
        reader.setOnImageAvailableListener({ r ->
            val image = r.acquireLatestImage() ?: return@setOnImageAvailableListener
            try {
                val plane = image.planes[0]
                val buffer = plane.buffer
                val rowStride = plane.rowStride
                val pixelStride = plane.pixelStride
                val rowBytes = width * pixelStride
                val expected = rowBytes * height
                if (rowStride == rowBytes && buffer.remaining() >= expected) {
                    buffer.get(frameBuffer, 0, expected)
                } else {
                    for (row in 0 until height) {
                        buffer.position(row * rowStride)
                        buffer.get(frameBuffer, row * rowBytes, rowBytes)
                    }
                }
                bridge.pushRootFrame(frameBuffer, width, height)
                framesDeliveredCounter.incrementAndGet()
            } catch (e: Exception) {
                Log.w(TAG, "Root frame delivery failed: ${e.message}")
            } finally {
                image.close()
            }
        }, handler)
        return reader
    }

    private fun buildDecoder(reader: ImageReader): MediaCodec {
        val format = MediaFormat.createVideoFormat(MIME_TYPE, width, height).apply {
            setInteger(MediaFormat.KEY_FRAME_RATE, fps)
        }
        val codec = MediaCodec.createDecoderByType(MIME_TYPE)
        codec.configure(format, reader.surface, null, 0)
        codec.start()
        return codec
    }

    private fun spawnScreenrecord(): Process? {
        // argv form — passes safely through Root.runAsRoot's shell-quote
        // logic so future changes to flag values can't introduce injection.
        val args = arrayOf(
            "screenrecord",
            "--output-format=h264",
            "--size=${width}x$height",
            "--bit-rate=$bitRate",
            // Time limit 0 isn't supported; the largest accepted is 180s.
            // We restart the process ourselves if it exits early.
            "--time-limit=180",
            "-",
        )
        // Inline ProcessBuilder so we have direct access to the child's
        // stdout (Root.runAsRoot returns Boolean). We still pass args
        // unquoted because the entire array is a fixed program+flags
        // with no user-controlled content.
        return try {
            ProcessBuilder("su", "-c", args.joinToString(" "))
                .redirectErrorStream(false)
                .start()
        } catch (e: Exception) {
            Log.e(TAG, "Failed to spawn `su -c screenrecord`: ${e.message}")
            null
        }
    }

    private fun startInputPump(initialStream: InputStream, codec: MediaCodec) {
        inputThread = Thread({
            val buf = ByteArray(INPUT_CHUNK)
            var stream: InputStream = initialStream
            try {
                outer@ while (running) {
                    val n = try {
                        stream.read(buf)
                    } catch (e: Exception) {
                        if (!running) break
                        Log.w(TAG, "screenrecord read error: ${e.message}")
                        -1
                    }
                    if (n <= 0) {
                        if (!running) break
                        // screenrecord caps at --time-limit=180s. When it
                        // exits cleanly we respawn so capture survives
                        // long sessions instead of freezing after ~3min.
                        Log.i(TAG, "screenrecord EOF — respawning")
                        synchronized(processLock) {
                            runCatching { process?.destroy() }
                            process = null
                        }
                        val next = spawnScreenrecord()
                        if (next == null) {
                            // Avoid a tight loop if `su` is suddenly unhappy.
                            try { Thread.sleep(500) } catch (_: InterruptedException) { break }
                            continue@outer
                        }
                        // Publish the new process under the lock so a
                        // concurrent stop() either (a) sees no process,
                        // tears down later, and lets us assign it for
                        // the destroy on the NEXT stop call — or (b) sees
                        // !running and we destroy the new process ourselves.
                        val accepted = synchronized(processLock) {
                            if (!running) {
                                false
                            } else {
                                process = next
                                true
                            }
                        }
                        if (!accepted) {
                            // running flipped false between EOF and now —
                            // someone called stop(). Drop the new process
                            // on the floor; the codec and output thread
                            // are stop()'s responsibility (it's the only
                            // writer to `running`, so we don't need to
                            // tear them down here).
                            runCatching { next.destroy() }
                            break@outer
                        }
                        stream = next.inputStream
                        continue@outer
                    }
                    var offset = 0
                    while (offset < n && running) {
                        val index = codec.dequeueInputBuffer(50_000)
                        if (index < 0) {
                            // Codec is starved — back off briefly instead
                            // of spinning. Without this, a stalled codec
                            // burns 100% of one core hammering dequeue.
                            try {
                                Thread.sleep(NO_BUFFER_BACKOFF_MS)
                            } catch (_: InterruptedException) {
                                break
                            }
                            continue
                        }
                        val inputBuffer = codec.getInputBuffer(index) ?: continue
                        inputBuffer.clear()
                        val chunk = minOf(n - offset, inputBuffer.capacity())
                        inputBuffer.put(buf, offset, chunk)
                        codec.queueInputBuffer(
                            index,
                            0,
                            chunk,
                            System.nanoTime() / 1_000,
                            0,
                        )
                        offset += chunk
                    }
                }
            } catch (_: InterruptedException) {
                // Expected on stop()
            } catch (e: Exception) {
                Log.w(TAG, "Input pump error: ${e.message}")
            }
        }, "RootDecoderInput").also { it.start() }
    }

    private fun startOutputDrain(codec: MediaCodec) {
        outputThread = Thread({
            val info = MediaCodec.BufferInfo()
            try {
                while (running) {
                    val index = codec.dequeueOutputBuffer(info, 50_000)
                    if (index >= 0) {
                        // `true` = render to the configured surface (ImageReader),
                        // which triggers onImageAvailable on the reader's handler.
                        codec.releaseOutputBuffer(index, true)
                        if (info.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM != 0) {
                            Log.i(TAG, "Decoder reported EOS")
                            break
                        }
                    }
                }
            } catch (_: InterruptedException) {
                // Expected on stop()
            } catch (e: Exception) {
                Log.w(TAG, "Output drain error: ${e.message}")
            }
        }, "RootDecoderOutput").also { it.start() }
    }
}
