/*
 * Copyright (C) 2025 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.eyedropper.util

import android.graphics.Bitmap
import android.hardware.HardwareBuffer
import android.os.IBinder
import android.util.Log
import android.view.IWindowManager
import android.window.IScreenCaptureCallback
import android.window.ScreenCapture
import javax.inject.Inject
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.withTimeout

private const val TAG = "ScreenCaptureHelper"

class ScreenCaptureHelper @Inject constructor(
    private val iWindowManager: IWindowManager,
) {
    /**
     * Most settings are default, but set to be explicit.
     */
    private fun getScreenCaptureParams(displayId: Int): ScreenCapture.ScreenCaptureParams {
        return ScreenCapture.ScreenCaptureParams.Builder(displayId)
            .setSecureContentPolicy(ScreenCapture.ScreenCaptureParams.SECURE_CONTENT_POLICY_REDACT)
            .setProtectedContentPolicy(
                ScreenCapture.ScreenCaptureParams.PROTECTED_CONTENT_POLICY_REDACT
            )
            .setPixelFormat(HardwareBuffer.RGBA_8888)
            .setUseDisplayInstallationOrientation(false)
            .setPreserveDisplayColors(false)
            .setCaptureMode(ScreenCapture.ScreenCaptureParams.CAPTURE_MODE_NONE)
            .setIncludeSystemOverlays(true)
            .build()
    }

    /**
     * Asynchronous wrapper around IWindowManager#screenCapture.
     */
    suspend fun captureScreenshotAsync(displayId: Int): Bitmap {
        val deferredBitmap = CompletableDeferred<Bitmap>()
        val params = getScreenCaptureParams(displayId)

        val callback = object : IScreenCaptureCallback.Stub() {

            override fun onSuccess(result: ScreenCapture.ScreenCaptureResult) {
                val hardwareBitmap = Bitmap.wrapHardwareBuffer(
                    result.hardwareBuffer, result.colorSpace)
                    ?: throw IllegalStateException("Bitmap.wrapHardwareBuffer returned null")
                Log.d(
                    TAG,
                    "Capture Success - HardwareBuffer: " +
                            "${result.hardwareBuffer.width}x${result.hardwareBuffer.height}"
                )

                // Can only getPixel() on software bitmap
                val softwareBitmap = hardwareBitmap.copy(Bitmap.Config.ARGB_8888, false)
                    ?: throw IllegalStateException(
                        "Error copying hardware bitmap to software bitmap")
                deferredBitmap.complete(softwareBitmap)
            }

            override fun onFailure(errorCode: Int) {
                Log.e(TAG, "Screen capture failure with errorCode == $errorCode")
                deferredBitmap.completeExceptionally(
                    RuntimeException("$ERROR_WITH_CODE $errorCode")
                )
            }

            override fun asBinder(): IBinder = this
        }

        return try {
            withTimeout(2000L) {
                iWindowManager.screenCapture(params, callback)
                deferredBitmap.await()
            }
        } catch (e: TimeoutCancellationException) {
            throw RuntimeException("$TIMED_OUT $e")
        } catch (e: Exception) {
            throw RuntimeException("$RUNTIME_ERROR $e")
        }
    }

    companion object {
        const val TIMED_OUT = "Screenshot capture timed out"
        const val RUNTIME_ERROR = "Screen capture failure with error code"
        const val ERROR_WITH_CODE = "Screenshot failed"
    }
}