package com.ahmed.jarvis

import android.app.Notification
import android.app.NotificationChannel
import android.app.NotificationManager
import android.app.Service
import android.content.Context
import android.content.Intent
import android.content.pm.ServiceInfo
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.media.projection.MediaProjectionManager
import android.os.Build
import android.os.Handler
import android.os.HandlerThread
import android.os.IBinder
import java.io.File
import java.io.FileOutputStream

/**
 * Captures a single screen frame via MediaProjection and saves it as a PNG.
 *
 * Android 14+ requires a foreground service of type `mediaProjection` to be
 * running BEFORE `MediaProjection.getMediaProjection` is called — hence the
 * whole capture lives here rather than in the Activity. The one-shot result
 * (saved path or error) is delivered back through [resultCallback], which
 * MainActivity sets right before starting the service.
 */
class ScreenCaptureService : Service() {

    companion object {
        const val EXTRA_CODE = "code"
        const val EXTRA_DATA = "data"
        const val EXTRA_WIDTH = "width"
        const val EXTRA_HEIGHT = "height"
        const val EXTRA_DPI = "dpi"

        private const val NOTIF_ID = 0xC0DE
        private const val CHANNEL_ID = "jarvis_capture"

        /** (ok, pathOrError). Set by MainActivity before startForegroundService. */
        @Volatile
        var resultCallback: ((Boolean, String) -> Unit)? = null
    }

    private var projection: MediaProjection? = null
    private var virtualDisplay: VirtualDisplay? = null
    private var imageReader: ImageReader? = null
    private var thread: HandlerThread? = null
    private var handled = false

    override fun onBind(intent: Intent?): IBinder? = null

    override fun onStartCommand(intent: Intent?, flags: Int, startId: Int): Int {
        startForegroundCompat()
        if (intent == null) {
            finish(false, "No capture request.")
            return START_NOT_STICKY
        }
        val code = intent.getIntExtra(EXTRA_CODE, 0)
        @Suppress("DEPRECATION")
        val data: Intent? = intent.getParcelableExtra(EXTRA_DATA)
        val width = intent.getIntExtra(EXTRA_WIDTH, 0)
        val height = intent.getIntExtra(EXTRA_HEIGHT, 0)
        val dpi = intent.getIntExtra(EXTRA_DPI, 0)
        if (data == null || width <= 0 || height <= 0) {
            finish(false, "Bad capture parameters.")
            return START_NOT_STICKY
        }

        try {
            val mpm = getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager
            val proj = mpm.getMediaProjection(code, data)
            if (proj == null) {
                finish(false, "Screen capture unavailable.")
                return START_NOT_STICKY
            }
            projection = proj

            val ht = HandlerThread("jarvis-capture").also { it.start() }
            thread = ht
            val handler = Handler(ht.looper)

            // Android 14 requires a registered callback before the virtual display.
            proj.registerCallback(object : MediaProjection.Callback() {}, handler)

            val reader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888, 2)
            imageReader = reader
            reader.setOnImageAvailableListener({ r ->
                if (handled) return@setOnImageAvailableListener
                val image = r.acquireLatestImage() ?: return@setOnImageAvailableListener
                try {
                    val path = saveImageAsPng(image, width, height)
                    finish(true, path)
                } catch (e: Exception) {
                    finish(false, e.message ?: "Save failed.")
                } finally {
                    image.close()
                }
            }, handler)

            virtualDisplay = proj.createVirtualDisplay(
                "jarvis-shot",
                width,
                height,
                dpi,
                DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
                reader.surface,
                null,
                handler
            )
        } catch (e: Exception) {
            finish(false, e.message ?: "Capture failed.")
        }
        return START_NOT_STICKY
    }

    private fun saveImageAsPng(image: android.media.Image, width: Int, height: Int): String {
        val plane = image.planes[0]
        val buffer = plane.buffer
        val pixelStride = plane.pixelStride
        val rowStride = plane.rowStride
        val rowPadding = rowStride - pixelStride * width
        val bmp = Bitmap.createBitmap(
            width + rowPadding / pixelStride,
            height,
            Bitmap.Config.ARGB_8888
        )
        bmp.copyPixelsFromBuffer(buffer)
        val cropped = if (rowPadding == 0) bmp else Bitmap.createBitmap(bmp, 0, 0, width, height)

        val dir = File(filesDir, "screenshots").apply { mkdirs() }
        val file = File(dir, "shot_${System.currentTimeMillis()}.png")
        FileOutputStream(file).use { out ->
            cropped.compress(Bitmap.CompressFormat.PNG, 100, out)
        }
        if (cropped !== bmp) cropped.recycle()
        bmp.recycle()
        return file.absolutePath
    }

    private fun finish(ok: Boolean, payload: String) {
        if (handled) return
        handled = true
        val cb = resultCallback
        resultCallback = null
        cleanup()
        cb?.invoke(ok, payload)
        stopForegroundCompat()
        stopSelf()
    }

    private fun cleanup() {
        try {
            virtualDisplay?.release()
        } catch (_: Exception) {}
        try {
            imageReader?.close()
        } catch (_: Exception) {}
        try {
            projection?.stop()
        } catch (_: Exception) {}
        try {
            thread?.quitSafely()
        } catch (_: Exception) {}
        virtualDisplay = null
        imageReader = null
        projection = null
        thread = null
    }

    private fun startForegroundCompat() {
        val nm = getSystemService(Context.NOTIFICATION_SERVICE) as NotificationManager
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            val ch = NotificationChannel(
                CHANNEL_ID,
                "Screen capture",
                NotificationManager.IMPORTANCE_LOW
            )
            nm.createNotificationChannel(ch)
        }
        val notif: Notification = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            Notification.Builder(this, CHANNEL_ID)
        } else {
            @Suppress("DEPRECATION")
            Notification.Builder(this)
        }
            .setContentTitle("Jarvis")
            .setContentText("Capturing the screen…")
            .setSmallIcon(android.R.drawable.ic_menu_camera)
            .setOngoing(true)
            .build()

        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
            startForeground(NOTIF_ID, notif, ServiceInfo.FOREGROUND_SERVICE_TYPE_MEDIA_PROJECTION)
        } else {
            startForeground(NOTIF_ID, notif)
        }
    }

    private fun stopForegroundCompat() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
            stopForeground(STOP_FOREGROUND_REMOVE)
        } else {
            @Suppress("DEPRECATION")
            stopForeground(true)
        }
    }

    override fun onDestroy() {
        cleanup()
        super.onDestroy()
    }
}