ScreenCaptureService.kt

package com.google.ai.sample

import android.app.Activity
import android.app.Notification
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.net.Uri
import android.os.Build
import android.os.Handler
import android.os.IBinder
import android.os.Looper
import android.util.DisplayMetrics
import android.util.Log
import android.view.WindowManager
import android.widget.Toast
import com.google.ai.client.generativeai.GenerativeModel
import com.google.ai.client.generativeai.type.Content
import com.google.ai.client.generativeai.type.ImagePart
import com.google.ai.client.generativeai.type.FunctionCallPart
import com.google.ai.client.generativeai.type.FunctionResponsePart
import com.google.ai.client.generativeai.type.BlobPart
import com.google.ai.client.generativeai.type.TextPart
import com.google.ai.sample.feature.multimodal.dtos.ContentDto
import com.google.ai.sample.feature.multimodal.dtos.toSdk
import com.google.ai.sample.service.AiCallRequestExtras
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.launch
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.json.Json
import kotlinx.serialization.decodeFromString
import kotlinx.serialization.MissingFieldException
import kotlinx.serialization.SerializationException
import kotlinx.serialization.json.JsonClassDiscriminator
import kotlinx.serialization.modules.SerializersModule
import kotlinx.serialization.modules.polymorphic
import kotlinx.serialization.modules.subclass
import androidx.localbroadcastmanager.content.LocalBroadcastManager
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import java.io.IOException

class ScreenCaptureService : Service() {
    private val screenCaptureStorage by lazy { ScreenCaptureStorage(applicationContext, TAG) }
    private val notificationFactory by lazy {
        ScreenCaptureNotificationFactory(applicationContext, CHANNEL_ID)
    }

    companion object {
        private const val TAG = "ScreenCaptureService"
        private const val CHANNEL_ID = "ScreenCaptureChannel"
        private const val NOTIFICATION_ID = 2001
        private const val NOTIFICATION_ID_AI = NOTIFICATION_ID + 1 // Or any distinct ID
        const val ACTION_START_CAPTURE = "com.google.ai.sample.START_CAPTURE"
        const val ACTION_TAKE_SCREENSHOT = "com.google.ai.sample.TAKE_SCREENSHOT" // New action
        const val ACTION_STOP_CAPTURE = "com.google.ai.sample.STOP_CAPTURE"   // New action
        const val ACTION_KEEP_ALIVE_FOR_WEBRTC = "com.google.ai.sample.KEEP_ALIVE_FOR_WEBRTC"
        const val EXTRA_RESULT_CODE = "result_code"
        const val EXTRA_RESULT_DATA = "result_data"
        const val EXTRA_TAKE_SCREENSHOT_ON_START = "take_screenshot_on_start"

        // For triggering AI call execution in the service
        const val ACTION_EXECUTE_AI_CALL = "com.google.ai.sample.EXECUTE_AI_CALL"
        const val EXTRA_AI_INPUT_CONTENT_JSON = "com.google.ai.sample.EXTRA_AI_INPUT_CONTENT_JSON"
        const val EXTRA_AI_CHAT_HISTORY_JSON = "com.google.ai.sample.EXTRA_AI_CHAT_HISTORY_JSON"
        const val EXTRA_AI_MODEL_NAME = "com.google.ai.sample.EXTRA_AI_MODEL_NAME" // For service to create model
        const val EXTRA_AI_API_KEY = "com.google.ai.sample.EXTRA_AI_API_KEY"     // For service to create model
        const val EXTRA_AI_API_PROVIDER = "com.google.ai.sample.EXTRA_AI_API_PROVIDER" // For service to select API
        const val EXTRA_TEMP_FILE_PATHS = "com.google.ai.sample.EXTRA_TEMP_FILE_PATHS"


        // For broadcasting AI call results from the service
        const val ACTION_AI_CALL_RESULT = "com.google.ai.sample.AI_CALL_RESULT"
        const val EXTRA_AI_RESPONSE_TEXT = "com.google.ai.sample.EXTRA_AI_RESPONSE_TEXT"
        const val EXTRA_AI_ERROR_MESSAGE = "com.google.ai.sample.EXTRA_AI_ERROR_MESSAGE"
        const val ACTION_AI_STREAM_UPDATE = "com.google.ai.sample.AI_STREAM_UPDATE"
        const val EXTRA_AI_STREAM_CHUNK = "com.google.ai.sample.EXTRA_AI_STREAM_CHUNK"

        private var instance: ScreenCaptureService? = null

        fun isRunning(): Boolean = instance != null && instance?.isReady == true
    }

    private var mediaProjection: MediaProjection? = null
    private var virtualDisplay: VirtualDisplay? = null
    private var imageReader: ImageReader? = null
    private var isReady = false // Flag to indicate if MediaProjection is set up and active
    private val isScreenshotRequestedRef = java.util.concurrent.atomic.AtomicBoolean(false)
    private val serviceScope = CoroutineScope(Dispatchers.IO + SupervisorJob())

    // Callback for MediaProjection
    private val mediaProjectionCallback = object : MediaProjection.Callback() {
        override fun onStop() {
            Log.w(TAG, "MediaProjection session stopped externally (via callback). Cleaning up.")
            cleanup() // Perform full cleanup if projection stops unexpectedly
        }
    }

    private fun startForegroundCompat(notificationId: Int, notification: Notification, foregroundType: Int) {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
            startForeground(notificationId, notification, foregroundType)
        } else {
            startForeground(notificationId, notification)
        }
    }

    private fun isHighDemandMessage(message: String?): Boolean {
        return message?.contains("503") == true ||
            message?.contains("overloaded") == true ||
            message?.contains("UNAVAILABLE") == true
    }

    private fun showLongToast(message: String) {
        Handler(Looper.getMainLooper()).post {
            Toast.makeText(applicationContext, message, Toast.LENGTH_LONG).show()
        }
    }

    private fun broadcastScreenshotCaptured(screenshotUri: Uri) {
        val intent = Intent(MainActivity.ACTION_MEDIAPROJECTION_SCREENSHOT_CAPTURED).apply {
            putExtra(MainActivity.EXTRA_SCREENSHOT_URI, screenshotUri.toString())
            `package` = applicationContext.packageName
        }
        applicationContext.sendBroadcast(intent)
        Log.d(TAG, "Sent broadcast ACTION_MEDIAPROJECTION_SCREENSHOT_CAPTURED with URI: $screenshotUri")
    }

    override fun onCreate() {
        super.onCreate()
        instance = this
        Log.d(TAG, "onCreate: Service created")
        createNotificationChannel()
    }

    override fun onStartCommand(intent: Intent?, flags: Int, startId: Int): Int {
        Log.d(TAG, "onStartCommand: action=${intent?.action}, isReady=$isReady, mediaProjectionIsNull=${mediaProjection==null}")

        when (intent?.action) {
            ACTION_KEEP_ALIVE_FOR_WEBRTC -> {
                Log.d(TAG, "Received ACTION_KEEP_ALIVE_FOR_WEBRTC. Starting foreground to hold MediaProjection token.")
                val notification = createNotification()
                startForegroundCompat(NOTIFICATION_ID, notification, ServiceInfo.FOREGROUND_SERVICE_TYPE_MEDIA_PROJECTION)
                isReady = true // Consider it ready so it doesn't try to start again
                return START_STICKY
            }
            ACTION_START_CAPTURE -> {
                if (isReady && mediaProjection != null) {
                    Log.w(TAG, "MediaProjection already active, ignoring duplicate START_CAPTURE")
                    return START_STICKY
                }
                val notification = createNotification()
                startForegroundCompat(NOTIFICATION_ID, notification, ServiceInfo.FOREGROUND_SERVICE_TYPE_MEDIA_PROJECTION)
                Log.d(TAG, "Service started in foreground for ACTION_START_CAPTURE.")

                val resultCode = intent.getIntExtra(EXTRA_RESULT_CODE, Activity.RESULT_CANCELED)
                val resultData = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) {
                    intent.getParcelableExtra(EXTRA_RESULT_DATA, Intent::class.java)
                } else {
                    @Suppress("DEPRECATION")
                    intent.getParcelableExtra<Intent>(EXTRA_RESULT_DATA)
                }

                Log.d(TAG, "onStartCommand (START_CAPTURE): resultCode=$resultCode, hasResultData=${resultData != null}")

                if (resultCode == Activity.RESULT_OK && resultData != null) {
                    val takeScreenshotFlag = intent.getBooleanExtra(EXTRA_TAKE_SCREENSHOT_ON_START, false)
                    startCapture(resultCode, resultData, takeScreenshotFlag)
                } else {
                    Log.e(TAG, "Invalid parameters for START_CAPTURE: resultCode=$resultCode (expected ${Activity.RESULT_OK}), resultDataIsNull=${resultData == null}")
                    cleanup() // Use cleanup to stop foreground and self
                }
            }
            ACTION_TAKE_SCREENSHOT -> {
                Log.d(TAG, "Received ACTION_TAKE_SCREENSHOT.")
                if (isReady && mediaProjection != null) {
                    takeScreenshot()
                } else {
                    Log.e(TAG, "Service not ready or MediaProjection not available for TAKE_SCREENSHOT. isReady=$isReady, mediaProjectionIsNull=${mediaProjection == null}")
                    showLongToast("Screenshot service not ready. Please re-grant permission if necessary.")
                    // Optionally, broadcast a failure or request MainActivity to re-initiate.
                    // If not ready, and this action is called, it implies a logic error or race condition.
                    // MainActivity should ideally prevent calling this if service isn't running/ready.
                }
            }
            ACTION_STOP_CAPTURE -> {
                Log.d(TAG, "Received ACTION_STOP_CAPTURE. Cleaning up.")
                cleanup()
            }
            ACTION_EXECUTE_AI_CALL -> {
                Log.d(TAG, "ACTION_EXECUTE_AI_CALL: Ensuring foreground state for AI processing.")
                val aiNotification = createAiOperationNotification()

                var startedForegroundForAi = false // Flag to track if we started foreground specifically for this call

                // Only start foreground if not already ready (i.e., not already in foreground with mediaProjection)
                if (!isReady) {
                    val foregroundType = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
                        ServiceInfo.FOREGROUND_SERVICE_TYPE_DATA_SYNC // Safer type for AI/network, no special permissions needed
                    } else {
                        0 // Use none for older versions
                    }
                    startForegroundCompat(NOTIFICATION_ID_AI, aiNotification, foregroundType)
                    Log.d(TAG, "Started foreground with type ${foregroundType} for AI processing (since not ready).")
                    startedForegroundForAi = true
                } else {
                    Log.d(TAG, "Already in foreground with mediaProjection, skipping startForeground for AI.")
                }

                Log.d(TAG, "Received ACTION_EXECUTE_AI_CALL")
                // This service, already a Foreground Service for MediaProjection,
                // is now also responsible for executing AI calls to leverage foreground network priority.
                val extras = AiCallRequestExtras.fromIntent(intent)
                val inputContentJson = extras.inputContentJson
                val chatHistoryJson = extras.chatHistoryJson
                val modelName = extras.modelName
                val apiKey = extras.apiKey
                val apiProvider = extras.apiProvider
                val tempFilePaths = extras.tempFilePaths
                Log.d(TAG, "Received tempFilePaths for cleanup: $tempFilePaths")

                if (inputContentJson == null || chatHistoryJson == null || modelName == null || apiKey == null) {
                    Log.e(TAG, "Missing necessary data for AI call. inputContentJson: ${inputContentJson != null}, chatHistoryJson: ${chatHistoryJson != null}, modelName: ${modelName != null}, apiKey: ${apiKey != null}")
                    // Optionally broadcast an error back immediately
                    broadcastAiCallError("Missing parameters for AI call in service.")
                    // If we started foreground for this, stop it now (but keep service running)
                    if (startedForegroundForAi) {
                        stopForeground(STOP_FOREGROUND_REMOVE)
                    }
                    return START_STICKY // Or START_NOT_STICKY if this is a fatal error for this call
                }

                serviceScope.launch {
                    var responseText: String? = null
                    var errorMessage: String? = null
                    try {
                        // Deserialize JSON to DTOs.
                        val chatHistoryDtos = Json.decodeFromString<List<ContentDto>>(chatHistoryJson)
                        val inputContentDto = Json.decodeFromString<ContentDto>(inputContentJson)

                        // Convert DTOs back to SDK types.
                        val chatHistory = chatHistoryDtos.map { it.toSdk() } // Uses ContentDto.toSdk()
                        val inputContent = inputContentDto.toSdk()           // Uses ContentDto.toSdk()

                        Log.d(TAG, "ACTION_EXECUTE_AI_CALL: Logging reloaded Bitmap properties after DTO conversion from file:")

                        // Log properties for inputContent's images
                        inputContent.parts.filterIsInstance<com.google.ai.client.generativeai.type.ImagePart>().forEachIndexed { index, imagePart ->
                            val bitmap = imagePart.image // This is the reloaded Bitmap
                            Log.d(TAG, "  InputContent Reloaded Image[${index}]: Width=${bitmap.width}, Height=${bitmap.height}, Config=${bitmap.config?.name ?: "null"}, HasAlpha=${bitmap.hasAlpha()}, IsMutable=${bitmap.isMutable}")
                        }

                        // Log properties for chat.history images
                        chatHistory.forEachIndexed { historyIndex, contentItem ->
                            contentItem.parts.filterIsInstance<com.google.ai.client.generativeai.type.ImagePart>().forEachIndexed { partIndex, imagePart ->
                                val bitmap = imagePart.image // This is the reloaded Bitmap
                                Log.d(TAG, "  History[${historyIndex}] Reloaded Image[${partIndex}]: Width=${bitmap.width}, Height=${bitmap.height}, Config=${bitmap.config?.name ?: "null"}, HasAlpha=${bitmap.hasAlpha()}, IsMutable=${bitmap.isMutable}")
                            }
                        }

                        Log.d(TAG, "ACTION_EXECUTE_AI_CALL: Saving reloaded Bitmaps for visual integrity check.")

                        // Save reloaded bitmaps from inputContent
                        inputContent.parts.filterIsInstance<com.google.ai.client.generativeai.type.ImagePart>().forEachIndexed { index, imagePart ->
                            val reloadedBitmap = imagePart.image
                            val reloadedBitmapDebugPath = com.google.ai.sample.util.ImageUtils.saveBitmapToTempFile(applicationContext, reloadedBitmap)
                            if (reloadedBitmapDebugPath != null) {
                                Log.d(TAG, "  InputContent Reloaded Image[${index}] (for debug) also saved to: $reloadedBitmapDebugPath. Compare with original.")
                            }
                        }

                        // Save reloaded bitmaps from chat.history
                        chatHistory.forEachIndexed { historyIndex, contentItem ->
                            contentItem.parts.filterIsInstance<com.google.ai.client.generativeai.type.ImagePart>().forEachIndexed { partIndex, imagePart ->
                                val reloadedBitmap = imagePart.image
                                val reloadedBitmapDebugPath = com.google.ai.sample.util.ImageUtils.saveBitmapToTempFile(applicationContext, reloadedBitmap)
                                if (reloadedBitmapDebugPath != null) {
                                    Log.d(TAG, "  History[${historyIndex}] Reloaded Image[${partIndex}] (for debug) also saved to: $reloadedBitmapDebugPath. Compare with original.")
                                }
                            }
                        }
                        try {
                            if (apiProvider == ApiProvider.VERCEL) {
                                responseText = callVercelApi(applicationContext, modelName, apiKey, chatHistoryDtos, inputContentDto)
                            } else if (apiProvider == ApiProvider.MISTRAL) {
                                val availableMistralKeys = ApiKeyManager.getInstance(applicationContext)
                                    .getApiKeys(ApiProvider.MISTRAL)
                                    .filter { it.isNotBlank() }
                                val result = callMistralApi(
                                    modelName = modelName,
                                    apiKey = apiKey,
                                    chatHistory = chatHistory,
                                    inputContent = inputContent,
                                    availableApiKeys = availableMistralKeys
                                )
                                responseText = result.first
                                errorMessage = result.second
                            } else if (apiProvider == ApiProvider.PUTER) {
                                val result = callPuterApi(modelName, apiKey, chatHistory, inputContent)
                                responseText = result.first
                                errorMessage = result.second
                            } else {
                                val generativeModel = GenerativeModel(
                                    modelName = modelName,
                                    apiKey = apiKey
                                )
                                val tempChat = generativeModel.startChat(history = chatHistory)
                                val fullResponse = StringBuilder()
                                tempChat.sendMessageStream(inputContent).collect { chunk ->
                                    chunk.text?.let {
                                        fullResponse.append(it)
                                        val streamIntent = Intent(ACTION_AI_STREAM_UPDATE).apply {
                                            putExtra(EXTRA_AI_STREAM_CHUNK, it)
                                        }
                                        LocalBroadcastManager.getInstance(applicationContext).sendBroadcast(streamIntent)
                                    }
                                }
                                responseText = fullResponse.toString()
                            }
                        } catch (e: MissingFieldException) {
                            Log.e(TAG, "Serialization error, potentially a 503 error.", e)
                            // Point 15: Check for missing 'parts' field (Gemma 27B issue)
                            if (e.message?.contains("parts") == true) {
                                errorMessage = "The model returned an incomplete response. This can happen with larger models. Please try again."
                            } else if (isHighDemandMessage(e.message)) {
                                // Point 14: User-friendly high-demand message
                                errorMessage = "This model is currently experiencing high demand. Please try again later."
                            } else {
                                errorMessage = e.localizedMessage ?: "Serialization error"
                            }
                        } catch (e: Exception) {
                            Log.e(TAG, "Direct error in AI call", e)
                            // Point 14: Check for high-demand 503 patterns
                            if (isHighDemandMessage(e.message)) {
                                errorMessage = "This model is currently experiencing high demand. Please try again later."
                            } else {
                                errorMessage = e.localizedMessage ?: "AI call failed"
                            }
                        }

                    } catch (e: Exception) {
                        // Catching general exceptions from model/chat operations or serialization
                        Log.e(TAG, "Outer error during AI call execution", e)

                        // Check if this is a 503-related error
                        if (e is MissingFieldException &&
                            (e.message?.contains("GRpcError") == true || isHighDemandMessage(e.message))) {
                            errorMessage = "This model is currently experiencing high demand. Please try again later."
                        } else if (e is MissingFieldException && e.message?.contains("parts") == true) {
                            // Point 15: Gemma 27B incomplete response
                            errorMessage = "The model returned an incomplete response. This can happen with larger models. Please try again."
                        } else if (isHighDemandMessage(e.message)) {
                            errorMessage = "This model is currently experiencing high demand. Please try again later."
                        } else {
                            errorMessage = e.localizedMessage ?: "Unknown error"
                        }
                    }
                    finally {
                        // Broadcast the result (success or error) back to the ViewModel.
                        val resultIntent = Intent(ACTION_AI_CALL_RESULT).apply {
                            if (responseText != null && errorMessage == null) {
                                putExtra(EXTRA_AI_RESPONSE_TEXT, responseText)
                            }
                            if (errorMessage != null) {
                                putExtra(EXTRA_AI_ERROR_MESSAGE, errorMessage)
                            }
                        }

                        if (errorMessage != null || responseText != null) {
                            LocalBroadcastManager.getInstance(applicationContext).sendBroadcast(resultIntent)
                            Log.d(TAG, "Local broadcast sent for AI_CALL_RESULT. Error: $errorMessage, Response: ${responseText != null}")
                        }

                        // Comment: Clean up temporary image files passed from the ViewModel.
                        if (tempFilePaths.isNotEmpty()) {
                            Log.d(TAG, "Cleaning up ${tempFilePaths.size} temporary image files.")
                            for (filePath in tempFilePaths) {
                                val deleted = com.google.ai.sample.util.ImageUtils.deleteFile(filePath)
                                if (!deleted) {
                                    Log.w(TAG, "Failed to delete temporary file: $filePath")
                                }
                            }
                        } else {
                            Log.d(TAG, "No temporary image files to clean up.")
                        }

                        // If we started foreground specifically for this AI call (i.e., !isReady), stop foreground now
                        // but KEEP THE SERVICE RUNNING (no stopSelf())
                        if (startedForegroundForAi) {
                            stopForeground(STOP_FOREGROUND_REMOVE)
                            Log.d(TAG, "Stopped foreground after AI call (since not ready), but service remains running.")
                        }
                    }
                }
                // START_STICKY to keep the service sticky/persistent
                return START_STICKY
            }
            else -> {
                Log.w(TAG, "Unknown or null action received: ${intent?.action}.")
                // If service is started with unknown action and not ready, stop it.
                if (!isReady) {
                    stopSelf()
                }
            }
        }
        return START_STICKY
    }

    private fun broadcastAiCallError(message: String) {
        val errorIntent = Intent(ACTION_AI_CALL_RESULT).apply {
            `package` = applicationContext.packageName
            putExtra(EXTRA_AI_ERROR_MESSAGE, message)
        }
        applicationContext.sendBroadcast(errorIntent)
        Log.d(TAG, "Broadcast error sent for AI_CALL_RESULT: $message")
    }

    private fun broadcastAiResult(responseText: String? = null, errorMessage: String? = null, isError: Boolean = false) {
        val resultIntent = Intent(ACTION_AI_CALL_RESULT).apply {
            if (responseText != null && !isError) {
                putExtra(EXTRA_AI_RESPONSE_TEXT, responseText)
            }
            if (errorMessage != null || isError) {
                putExtra(EXTRA_AI_ERROR_MESSAGE, errorMessage ?: "An unknown error occurred.")
            }
        }
        LocalBroadcastManager.getInstance(applicationContext).sendBroadcast(resultIntent)
        Log.d(TAG, "Local broadcast sent for AI_CALL_RESULT. Error: $errorMessage, Response: ${responseText != null}")
    }

    private fun createAiOperationNotification(): Notification {
        return notificationFactory.createAiOperationNotification()
    }

    private fun createNotification(): Notification {
        return notificationFactory.createNotification()
    }

    private fun createNotificationChannel() {
        notificationFactory.createNotificationChannel("Screen Capture Service")
        Log.d(TAG, "Notification channel created")
    }

    private fun startCapture(resultCode: Int, data: Intent, takeScreenshotOnStart: Boolean) {
        try {
            Log.d(TAG, "startCapture: Getting MediaProjection, takeScreenshotOnStart: $takeScreenshotOnStart")
            val mediaProjectionManager = getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager

            mediaProjection?.unregisterCallback(mediaProjectionCallback) // Unregister old before stopping
            mediaProjection?.stop() // Stop any existing projection

            mediaProjection = mediaProjectionManager.getMediaProjection(resultCode, data)

            if (mediaProjection == null) {
                Log.e(TAG, "MediaProjection is null after getMediaProjection call")
                isReady = false
                cleanup() // Use cleanup to stop foreground and self
                return
            }
            mediaProjection?.registerCallback(mediaProjectionCallback, Handler(Looper.getMainLooper()))
            isReady = true
            Log.d(TAG, "MediaProjection ready.")

            if (takeScreenshotOnStart) {
                Handler(Looper.getMainLooper()).postDelayed({
                    if(isReady && mediaProjection != null) {
                        Log.d(TAG, "startCapture: Taking initial screenshot after delay because takeScreenshotOnStart was true.")
                        takeScreenshot()
                    } else {
                        Log.w(TAG, "startCapture: Conditions to take initial screenshot not met after delay, even though takeScreenshotOnStart was true. isReady=$isReady, mediaProjectionIsNull=${mediaProjection==null}")
                    }
                }, 500)
            } else {
                Log.d(TAG, "startCapture: MediaProjection initialized, but skipping immediate screenshot as takeScreenshotOnStart is false.")
            }

        } catch (e: Exception) {
            Log.e(TAG, "Error in startCapture", e)
            isReady = false
            cleanup() // Use cleanup to stop foreground and self
        }
    }

    private fun takeScreenshot() {
        if (!isReady || mediaProjection == null) {
            Log.e(TAG, "Cannot take screenshot - service not ready or mediaProjection is null. isReady=$isReady, mediaProjectionIsNull=${mediaProjection == null}")
            return
        }
        isScreenshotRequestedRef.set(true)
        Log.d(TAG, "takeScreenshot: Preparing to capture. isScreenshotRequestedRef set to true.")

        try {
            // Check if we need to initialize VirtualDisplay and ImageReader
            if (virtualDisplay == null || imageReader == null) {
                val windowManager = getSystemService(Context.WINDOW_SERVICE) as WindowManager
                val displayMetrics = DisplayMetrics()

                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                    val bounds = windowManager.currentWindowMetrics.bounds
                    displayMetrics.widthPixels = bounds.width()
                    displayMetrics.heightPixels = bounds.height()
                    displayMetrics.densityDpi = resources.displayMetrics.densityDpi
                } else {
                    @Suppress("DEPRECATION")
                    windowManager.defaultDisplay.getMetrics(displayMetrics)
                }

                val width = displayMetrics.widthPixels
                val height = displayMetrics.heightPixels
                val density = displayMetrics.densityDpi

                if (width <= 0 || height <= 0) {
                    Log.e(TAG, "Invalid display dimensions: ${width}x${height}. Cannot create ImageReader.")
                    return
                }
                Log.d(TAG, "Display dimensions: ${width}x${height}, density: $density")

                imageReader?.close() // Close previous reader if any
                virtualDisplay?.release() // Release previous display if any

                imageReader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888, 1)
                val localImageReader = imageReader ?: run {
                    Log.e(TAG, "ImageReader is null after creation attempt.")
                    return
                }

                localImageReader.setOnImageAvailableListener({ reader ->
                    if (isScreenshotRequestedRef.compareAndSet(true, false)) {
                        Log.d(TAG, "Screenshot request flag consumed, processing image.")
                        var image: android.media.Image? = null
                        try {
                            image = reader.acquireLatestImage()
                            if (image != null) {
                                val planes = image.planes
                                val buffer = planes[0].buffer
                                val pixelStride = planes[0].pixelStride
                                val rowStride = planes[0].rowStride
                                val rowPadding = rowStride - pixelStride * width

                                val bitmap = Bitmap.createBitmap(
                                    width + rowPadding / pixelStride,
                                    height,
                                    Bitmap.Config.ARGB_8888
                                )
                                bitmap.copyPixelsFromBuffer(buffer)
                                Log.d(TAG, "Bitmap created, proceeding to save.")
                                saveScreenshot(bitmap)
                            } else {
                                Log.w(TAG, "acquireLatestImage returned null despite requested flag.")
                            }
                        } catch (e: IllegalStateException) {
                            Log.e(TAG, "Error processing image in listener", e)
                        } finally {
                            image?.close()
                            // Do NOT release VirtualDisplay or ImageReader here
                            // They will be reused for the next screenshot
                            Log.d(TAG, "Screenshot processed (or attempted), keeping resources for reuse.")
                        }
                    } else {
                        // Logic to discard the frame if no screenshot was formally requested
                        var imageToDiscard: android.media.Image? = null
                        try {
                            imageToDiscard = reader.acquireLatestImage()
                        } catch (e: IllegalStateException) {
                            // This catch is important because acquireLatestImage can fail if buffers are truly messed up
                            Log.e(TAG, "Error acquiring image to discard in OnImageAvailableListener else block", e)
                        } finally {
                            imageToDiscard?.close()
                        }
                    }
                }, Handler(Looper.getMainLooper()))

                virtualDisplay = mediaProjection?.createVirtualDisplay(
                    "ScreenCapture",
                    width, height, density,
                    DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
                    localImageReader.surface,
                    object : VirtualDisplay.Callback() {
                        override fun onPaused() { Log.d(TAG, "VirtualDisplay paused") }
                        override fun onResumed() { Log.d(TAG, "VirtualDisplay resumed") }
                        override fun onStopped() { Log.d(TAG, "VirtualDisplay stopped") }
                    },
                    Handler(Looper.getMainLooper())
                )

                if (virtualDisplay == null) {
                    Log.e(TAG, "Failed to create VirtualDisplay.")
                    localImageReader.close() // Clean up the reader we just created
                    this.imageReader = null
                    return
                }
                Log.d(TAG, "VirtualDisplay and ImageReader initialized for reuse.")
            } else {
                // Resources already exist, just trigger a new capture
                Log.d(TAG, "Using existing VirtualDisplay and ImageReader.")
                // Force the ImageReader to capture a new frame
                // The listener is already set up and will handle the new image
            }

        } catch (e: IllegalStateException) {
            Log.e(TAG, "Error in takeScreenshot setup", e)
            virtualDisplay?.release()
            virtualDisplay = null
            imageReader?.close()
            imageReader = null
        }
    }

    private fun saveScreenshot(bitmap: Bitmap) {
        screenCaptureStorage.saveScreenshot(
            bitmap = bitmap,
            onSaved = { screenshotUri -> broadcastScreenshotCaptured(screenshotUri) },
            onSuccessMessage = { message -> showLongToast(message) },
            onErrorMessage = { message -> showLongToast(message) }
        )
    }

    private fun cleanup() {
        Log.d(TAG, "cleanup() called. Cleaning up all MediaProjection resources.")
        try {
            isReady = false
            virtualDisplay?.release()
            virtualDisplay = null
            imageReader?.close()
            imageReader = null

            mediaProjection?.unregisterCallback(mediaProjectionCallback)
            mediaProjection?.stop()
            mediaProjection = null
        } catch (e: IllegalStateException) {
            Log.e(TAG, "Error during full cleanup", e)
        } finally {
            stopForeground(STOP_FOREGROUND_REMOVE)
            stopSelf() // This will trigger onDestroy eventually
            instance = null // Clear static instance
            Log.d(TAG, "Full cleanup finished, service fully stopped.")
        }
    }

    override fun onDestroy() {
        Log.d(TAG, "onDestroy: Service being destroyed")
        // Cleanup is called from ACTION_STOP_CAPTURE or if projection stops externally.
        // If service is killed by system, this ensures cleanup too.
        if (isReady || mediaProjection != null) { // Check if cleanup is actually needed
           cleanup()
        }
        serviceScope.cancel() // Cancel all coroutines in this scope
        instance = null // Ensure instance is cleared
        super.onDestroy()
    }

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