viewcapturelib/src/com/android/app/viewcapture/ViewCaptureAwareWindowManager.kt

/*
 * Copyright (C) 2024 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.app.viewcapture

import android.media.permission.SafeCloseable
import android.view.View
import android.view.ViewGroup
import android.view.WindowManager

/** Tag for debug logging. */
private const val TAG = "ViewCaptureWindowManager"

/**
 * Wrapper class for [WindowManager]. Adds [ViewCapture] to associated window when it is added to
 * view hierarchy.
 */
class ViewCaptureAwareWindowManager(
    private val windowManager: WindowManager,
    private val lazyViewCapture: Lazy<ViewCapture>,
    private val isViewCaptureEnabled: Boolean,
) : WindowManager by windowManager {

    private var viewCaptureCloseableMap: MutableMap<View, SafeCloseable> = mutableMapOf()

    override fun addView(view: View, params: ViewGroup.LayoutParams?) {
        windowManager.addView(view, params)
        if (isViewCaptureEnabled) {
            val viewCaptureCloseable: SafeCloseable =
                lazyViewCapture.value.startCapture(view, getViewName(view))
            viewCaptureCloseableMap[view] = viewCaptureCloseable
        }
    }

    override fun removeView(view: View?) {
        removeViewFromCloseableMap(view)
        windowManager.removeView(view)
    }

    override fun removeViewImmediate(view: View?) {
        removeViewFromCloseableMap(view)
        windowManager.removeViewImmediate(view)
    }

    private fun getViewName(view: View) = "." + view.javaClass.name

    private fun removeViewFromCloseableMap(view: View?) {
        if (isViewCaptureEnabled) {
            if (viewCaptureCloseableMap.containsKey(view)) {
                viewCaptureCloseableMap[view]?.close()
                viewCaptureCloseableMap.remove(view)
            }
        }
    }

    interface Factory {
        fun create(windowManager: WindowManager): ViewCaptureAwareWindowManager
    }
}