/*
 * Copyright (C) 2022 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.launcher3.statehandlers

import android.content.Context
import android.util.Log
import android.util.SparseArray
import android.util.SparseBooleanArray
import android.view.Display.DEFAULT_DISPLAY
import androidx.core.util.set
import com.android.internal.util.LatencyTracker
import com.android.launcher3.LauncherState
import com.android.launcher3.dagger.ApplicationContext
import com.android.launcher3.dagger.LauncherAppComponent
import com.android.launcher3.dagger.LauncherAppSingleton
import com.android.launcher3.statemanager.BaseState
import com.android.launcher3.util.DaggerSingletonObject
import com.android.launcher3.util.DaggerSingletonTracker
import com.android.launcher3.util.Executors.MAIN_EXECUTOR
import com.android.launcher3.util.window.WindowManagerProxy.DesktopVisibilityListener
import com.android.quickstep.SystemUiProxy
import com.android.quickstep.fallback.RecentsState
import com.android.wm.shell.desktopmode.DisplayDeskState
import com.android.wm.shell.desktopmode.IDesktopTaskListener.Stub
import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.enableMultipleDesktops
import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.useRoundedCorners
import java.io.PrintWriter
import java.lang.ref.WeakReference
import javax.inject.Inject

/**
 * Controls the visibility of the workspace and the resumed / paused state when desktop mode is
 * enabled.
 */
@LauncherAppSingleton
class DesktopVisibilityController
@Inject
constructor(
    @ApplicationContext private val context: Context,
    systemUiProxy: SystemUiProxy,
    lifecycleTracker: DaggerSingletonTracker,
) {
    /**
     * Tracks the desks configurations on each display.
     *
     * (Used only when multiple desks are enabled).
     *
     * @property displayId The ID of the display this object represents.
     * @property activeDeskId The ID of the active desk on the associated display (if any). It has a
     *   value of `INACTIVE_DESK_ID` (-1) if there are no active desks. Note that there can only be
     *   at most one active desk on each display.
     * @property deskIds a set containing the IDs of the desks on the associated display.
     */
    private data class DisplayDeskConfig(
        val displayId: Int,
        var activeDeskId: Int = INACTIVE_DESK_ID,
        val deskIds: MutableSet<Int>,
    )

    /** True if it is possible to create new desks on current setup. */
    var canCreateDesks: Boolean = false
        private set(value) {
            if (field == value) return
            field = value
            desktopVisibilityListeners.forEach { it.onCanCreateDesksChanged(field) }
        }

    /** Maps each display by its ID to its desks configuration. */
    private val displaysDesksConfigsMap = SparseArray<DisplayDeskConfig>()

    private val desktopVisibilityListeners: MutableSet<DesktopVisibilityListener> = HashSet()
    private val taskbarDesktopModeListeners: MutableSet<TaskbarDesktopModeListener> = HashSet()

    // This simply indicates that user is currently in desktop mode or not.
    @Deprecated("Does not work with multi-desks") private var isInDesktopModeDeprecated = false

    // to let launcher hold off on notifying desktop visibility listeners.
    var launcherAnimationRunning = false

    // TODO: b/394387739 - Deprecate this and replace it with something that tracks the count per
    //  desk.
    /**
     * Number of visible desktop windows in desktop mode. This can be > 0 when user goes to overview
     * from desktop window mode.
     */
    @Deprecated("Does not work with multi-desks")
    var visibleDesktopTasksCountDeprecated: Int = 0
        /**
         * Sets the number of desktop windows that are visible and updates launcher visibility based
         * on it.
         */
        set(visibleTasksCount) {
            if (enableMultipleDesktops(context)) {
                return
            }
            if (DEBUG) {
                Log.d(
                    TAG,
                    ("setVisibleDesktopTasksCount: visibleTasksCount=" +
                        visibleTasksCount +
                        " currentValue=" +
                        field),
                )
            }

            if (visibleTasksCount != field) {
                if (visibleDesktopTasksCountDeprecated == 0 && visibleTasksCount == 1) {
                    isInDesktopModeDeprecated = true
                }
                if (visibleDesktopTasksCountDeprecated == 1 && visibleTasksCount == 0) {
                    isInDesktopModeDeprecated = false
                }
            }
        }

    private var inOverviewStateMap = SparseBooleanArray()

    private var desktopTaskListener: DesktopTaskListenerImpl?

    init {
        desktopTaskListener = DesktopTaskListenerImpl(this, context)
        systemUiProxy.setDesktopTaskListener(desktopTaskListener)

        lifecycleTracker.addCloseable {
            desktopTaskListener = null
            systemUiProxy.setDesktopTaskListener(null)
        }
    }

    /**
     * Returns the ID of the active desk (if any) on the display whose ID is [displayId], or
     * [INACTIVE_DESK_ID] if no desk is currently active or the multiple desks feature is disabled.
     */
    fun getActiveDeskId(displayId: Int): Int {
        if (!enableMultipleDesktops(context)) {
            // When the multiple desks feature is disabled, callers should not rely on the concept
            // of a desk ID.
            return INACTIVE_DESK_ID
        }

        return getDisplayDeskConfig(displayId)?.activeDeskId ?: INACTIVE_DESK_ID
    }

    /** Returns whether a desk is currently active on the display with the given [displayId]. */
    fun isInDesktopMode(displayId: Int): Boolean {
        if (!enableMultipleDesktops(context)) {
            return isInDesktopModeDeprecated
        }

        val activeDeskId = getDisplayDeskConfig(displayId)?.activeDeskId ?: INACTIVE_DESK_ID
        val isInDesktopMode = activeDeskId != INACTIVE_DESK_ID
        if (DEBUG) {
            Log.d(TAG, "isInDesktopMode ($displayId): $isInDesktopMode")
        }
        return isInDesktopMode
    }

    /**
     * Returns whether a desk is currently active on the display with the given [displayId] and
     * Overview is not active.
     */
    fun isInDesktopModeAndNotInOverview(displayId: Int): Boolean {
        if (!enableMultipleDesktops(context)) {
            return areDesktopTasksVisibleAndNotInOverview(displayId)
        }
        val inOverviewState = inOverviewStateMap[displayId]
        if (DEBUG) {
            Log.d(
                TAG,
                "isInDesktopModeAndNotInOverview: displayId=$displayId overview=$inOverviewState",
            )
        }
        return isInDesktopMode(displayId) && !inOverviewState
    }

    /** Whether desktop tasks are visible in desktop mode. */
    private fun areDesktopTasksVisibleAndNotInOverview(displayId: Int): Boolean {
        val desktopTasksVisible: Boolean = visibleDesktopTasksCountDeprecated > 0
        val inOverviewState = inOverviewStateMap[displayId]
        if (DEBUG) {
            Log.d(
                TAG,
                ("areDesktopTasksVisible: displayId=$displayId desktopVisible=$desktopTasksVisible" +
                    " overview=$inOverviewState"),
            )
        }
        return desktopTasksVisible && !inOverviewState
    }

    /** Registers a listener for Taskbar changes in Desktop Mode. */
    fun registerTaskbarDesktopModeListener(listener: TaskbarDesktopModeListener) {
        taskbarDesktopModeListeners.add(listener)
    }

    /** Removes a previously registered listener for Taskbar changes in Desktop Mode. */
    fun unregisterTaskbarDesktopModeListener(listener: TaskbarDesktopModeListener) {
        taskbarDesktopModeListeners.remove(listener)
    }

    fun onLauncherStateChanged(displayId: Int, state: LauncherState) {
        onLauncherStateChanged(
            state,
            state === LauncherState.BACKGROUND_APP,
            state.isRecentsViewVisible,
            displayId,
        )
    }

    /**
     * Launcher Driven Desktop Mode changes. For example, swipe to home and quick switch from
     * Desktop Windowing Mode. if there is any pending notification please notify desktop visibility
     * listeners.
     */
    fun onLauncherAnimationFromDesktopEnd(displayId: Int) {
        launcherAnimationRunning = false
    }

    fun onLauncherStateChanged(displayId: Int, state: RecentsState) {
        onLauncherStateChanged(
            state,
            state === RecentsState.BACKGROUND_APP,
            state.isRecentsViewVisible(),
            displayId,
        )
    }

    /** Process launcher state change and update launcher view visibility based on desktop state */
    fun onLauncherStateChanged(
        state: BaseState<*>,
        isBackgroundAppState: Boolean,
        isRecentsViewVisible: Boolean,
        displayId: Int,
    ) {
        if (DEBUG) {
            Log.d(TAG, "onLauncherStateChanged: newState=$state")
        }
        // Desktop visibility tracks overview and background state separately
        setOverviewStateEnabled(displayId, !isBackgroundAppState && isRecentsViewVisible)
    }

    private fun setOverviewStateEnabled(displayId: Int, overviewStateEnabled: Boolean) {
        val inOverviewState = inOverviewStateMap[displayId]
        if (DEBUG) {
            Log.d(
                TAG,
                ("setOverviewStateEnabled: enabled=" +
                    overviewStateEnabled +
                    " currentValue=" +
                    inOverviewState),
            )
        }
        if (overviewStateEnabled != inOverviewState) {
            inOverviewStateMap[displayId] = overviewStateEnabled
        }
    }

    /** Registers a listener for Taskbar changes in Desktop Mode. */
    fun registerDesktopVisibilityListener(listener: DesktopVisibilityListener) {
        desktopVisibilityListeners.add(listener)
    }

    /** Removes a previously registered listener for Taskbar changes in Desktop Mode. */
    fun unregisterDesktopVisibilityListener(listener: DesktopVisibilityListener) {
        desktopVisibilityListeners.remove(listener)
    }

    private fun notifyTaskbarDesktopModeListeners(doesAnyTaskRequireTaskbarRounding: Boolean) {
        if (DEBUG) {
            Log.d(
                TAG,
                "notifyTaskbarDesktopModeListeners: doesAnyTaskRequireTaskbarRounding=" +
                    doesAnyTaskRequireTaskbarRounding,
            )
        }
        for (listener in taskbarDesktopModeListeners) {
            listener.onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding)
        }
    }

    private fun notifyTaskbarDesktopModeListenersForEntry(displayId: Int, duration: Int) {
        if (DEBUG) {
            Log.d(TAG, "notifyTaskbarDesktopModeListenersForEntry: duration=" + duration)
        }
        for (listener in taskbarDesktopModeListeners) {
            listener.onEnterDesktopMode(duration)
        }
    }

    private fun notifyTaskbarDesktopModeListenersForExit(displayId: Int, duration: Int) {
        if (DEBUG) {
            Log.d(TAG, "notifyTaskbarDesktopModeListenersForExit: duration=" + duration)
        }
        for (listener in taskbarDesktopModeListeners) {
            listener.onExitDesktopMode(duration)
        }
    }

    private fun notifyOnDeskAdded(displayId: Int, deskId: Int) {
        if (DEBUG) {
            Log.d(TAG, "notifyOnDeskAdded: displayId=$displayId, deskId=$deskId")
        }

        for (listener in desktopVisibilityListeners) {
            listener.onDeskAdded(displayId, deskId)
        }
    }

    private fun notifyOnDeskRemoved(displayId: Int, deskId: Int) {
        if (DEBUG) {
            Log.d(TAG, "notifyOnDeskRemoved: displayId=$displayId, deskId=$deskId")
        }

        for (listener in desktopVisibilityListeners) {
            listener.onDeskRemoved(displayId, deskId)
        }
    }

    private fun notifyOnActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {
        if (DEBUG) {
            Log.d(
                TAG,
                "notifyOnActiveDeskChanged: displayId=$displayId, newActiveDesk=$newActiveDesk, oldActiveDesk=$oldActiveDesk",
            )
        }

        for (listener in desktopVisibilityListeners) {
            listener.onActiveDeskChanged(displayId, newActiveDesk, oldActiveDesk)
        }
    }

    // Called when the DesktopTaskListener is first connected to WM.
    private fun onListenerConnected(
        displayDeskStates: Array<DisplayDeskState>,
        canCreateDesks: Boolean,
    ) {
        if (!enableMultipleDesktops(context)) {
            return
        }

        displaysDesksConfigsMap.clear()

        displayDeskStates.forEach { displayDeskState ->
            if (DEBUG) {
                Log.d(TAG, "onListenerConnected displayId=${displayDeskState.displayId}")
            }
            displaysDesksConfigsMap[displayDeskState.displayId] =
                DisplayDeskConfig(
                    displayId = displayDeskState.displayId,
                    activeDeskId = displayDeskState.activeDeskId,
                    deskIds = displayDeskState.deskIds.toMutableSet(),
                )
        }

        this.canCreateDesks = canCreateDesks

        notifyOnListenerInitializedFromShell()
    }

    private fun notifyOnListenerInitializedFromShell() {
        if (DEBUG) {
            Log.d(TAG, "notifyOnListenerInitializedFromShell")
        }

        for (listener in desktopVisibilityListeners) {
            listener.onListenerInitializedFromShell()
        }
    }

    private fun getDisplayDeskConfig(displayId: Int) = displaysDesksConfigsMap[displayId]

    private fun onCanCreateDesksChanged(canCreateDesks: Boolean) {
        if (!enableMultipleDesktops(context)) {
            return
        }

        this.canCreateDesks = canCreateDesks
    }

    private fun onDeskAdded(displayId: Int, deskId: Int) {
        if (!enableMultipleDesktops(context)) {
            return
        }

        // Add the config for the desk if there is nothing yet, as the display can start without any
        // desks.
        if (getDisplayDeskConfig(displayId) == null) {
            displaysDesksConfigsMap[displayId] =
                DisplayDeskConfig(displayId, INACTIVE_DESK_ID, mutableSetOf(deskId))
        } else {
            getDisplayDeskConfig(displayId)!!.also {
                check(it.deskIds.add(deskId)) {
                    "Found a duplicate desk Id: $deskId on display: $displayId"
                }
            }
        }

        notifyOnDeskAdded(displayId, deskId)
    }

    private fun onDeskRemoved(displayId: Int, deskId: Int) {
        if (!enableMultipleDesktops(context)) {
            return
        }

        getDisplayDeskConfig(displayId)?.also {
            check(it.deskIds.remove(deskId)) {
                "Removing non-existing desk Id: $deskId on display: $displayId"
            }
            if (it.activeDeskId == deskId) {
                it.activeDeskId = INACTIVE_DESK_ID
            }
        }

        notifyOnDeskRemoved(displayId, deskId)
    }

    private fun onActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {
        if (!enableMultipleDesktops(context)) {
            return
        }

        getDisplayDeskConfig(displayId)?.also {
            check(oldActiveDesk == it.activeDeskId) {
                "Mismatch between the Shell's oldActiveDesk: $oldActiveDesk, and Launcher's: ${it.activeDeskId}"
            }
            check(newActiveDesk == INACTIVE_DESK_ID || it.deskIds.contains(newActiveDesk)) {
                "newActiveDesk: $newActiveDesk was never added to display: $displayId"
            }
            it.activeDeskId = newActiveDesk
        }

        if (newActiveDesk != oldActiveDesk) {
            notifyOnActiveDeskChanged(displayId, newActiveDesk, oldActiveDesk)
        }
    }

    fun dumpLogs(prefix: String, pw: PrintWriter) {
        pw.println(prefix + "DesktopVisibilityController:")

        pw.println("$prefix\tdesktopVisibilityListeners=$desktopVisibilityListeners")
        pw.println("$prefix\tvisibleDesktopTasksCount=$visibleDesktopTasksCountDeprecated")
        pw.println("$prefix\tinOverviewState=$inOverviewStateMap")
        pw.println("$prefix\tdesktopTaskListener=$desktopTaskListener")
        pw.println("$prefix\tcontext=$context")
    }

    /**
     * Wrapper for the IDesktopTaskListener stub to prevent lingering references to the launcher
     * activity via the controller.
     */
    private class DesktopTaskListenerImpl(
        controller: DesktopVisibilityController,
        @ApplicationContext private val context: Context,
    ) : Stub() {
        private val controller = WeakReference(controller)
        private val displayId = context.displayId

        override fun onListenerConnected(
            displayDeskStates: Array<DisplayDeskState>,
            canCreateDesks: Boolean,
        ) {
            MAIN_EXECUTOR.execute {
                controller.get()?.onListenerConnected(displayDeskStates, canCreateDesks)
            }
        }

        @Deprecated("Not needed by multi-desks")
        override fun onTasksVisibilityChanged(displayId: Int, visibleTasksCount: Int) {
            if (enableMultipleDesktops(context)) return
            if (displayId != this.displayId) return
            MAIN_EXECUTOR.execute {
                controller.get()?.apply {
                    if (DEBUG) {
                        Log.d(TAG, "desktop visible tasks count changed=$visibleTasksCount")
                    }
                    visibleDesktopTasksCountDeprecated = visibleTasksCount
                }
            }
        }

        override fun onStashedChanged(displayId: Int, stashed: Boolean) {
            Log.w(TAG, "DesktopTaskListenerImpl: onStashedChanged is deprecated")
        }

        override fun onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding: Boolean) {
            if (!useRoundedCorners()) return
            MAIN_EXECUTOR.execute {
                controller.get()?.apply {
                    Log.d(
                        TAG,
                        "DesktopTaskListenerImpl: doesAnyTaskRequireTaskbarRounding= " +
                            doesAnyTaskRequireTaskbarRounding,
                    )
                    notifyTaskbarDesktopModeListeners(doesAnyTaskRequireTaskbarRounding)
                }
            }
        }

        @Deprecated("Not needed by multi-desks")
        override fun onEnterDesktopModeTransitionStarted(transitionDuration: Int) {
            if (enableMultipleDesktops(context)) return
            val controller = controller.get() ?: return
            MAIN_EXECUTOR.execute {
                Log.d(
                    TAG,
                    ("DesktopTaskListenerImpl: onEnterDesktopModeTransitionStarted with " +
                        "duration= " +
                        transitionDuration),
                )
                if (!controller.isInDesktopModeDeprecated) {
                    controller.isInDesktopModeDeprecated = true
                    controller.notifyTaskbarDesktopModeListenersForEntry(
                        DEFAULT_DISPLAY,
                        transitionDuration,
                    )
                }
            }
        }

        @Deprecated("Not needed by multi-desks")
        override fun onExitDesktopModeTransitionStarted(
            transitionDuration: Int,
            shouldEndUpAtHome: Boolean,
        ) {
            if (enableMultipleDesktops(context)) return
            val controller = controller.get() ?: return
            MAIN_EXECUTOR.execute {
                Log.d(
                    TAG,
                    ("DesktopTaskListenerImpl: onExitDesktopModeTransitionStarted with " +
                        "duration= " +
                        transitionDuration),
                )
                // If shouldEndUpAtHome is true, desktop mode is ending from the user
                // closing/minimizing the last open window. If it's false, the display is
                // probably transitioning to an app's full screen mode instead so this metric
                // should not be logged.
                if (shouldEndUpAtHome) {
                    LatencyTracker.getInstance(context)
                        .onActionStart(
                            LatencyTracker.ACTION_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE
                        )
                }
                if (controller.isInDesktopModeDeprecated) {
                    controller.isInDesktopModeDeprecated = false
                    controller.notifyTaskbarDesktopModeListenersForExit(
                        DEFAULT_DISPLAY,
                        transitionDuration,
                    )
                }
            }
        }

        override fun onCanCreateDesksChanged(canCreateDesks: Boolean) {
            MAIN_EXECUTOR.execute { controller.get()?.onCanCreateDesksChanged(canCreateDesks) }
        }

        override fun onDeskAdded(displayId: Int, deskId: Int) {
            MAIN_EXECUTOR.execute { controller.get()?.onDeskAdded(displayId, deskId) }
        }

        override fun onDeskRemoved(displayId: Int, deskId: Int) {
            MAIN_EXECUTOR.execute { controller.get()?.onDeskRemoved(displayId, deskId) }
        }

        override fun onActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {
            MAIN_EXECUTOR.execute {
                controller.get()?.onActiveDeskChanged(displayId, newActiveDesk, oldActiveDesk)
            }
        }
    }

    /** A listener for Taskbar in Desktop Mode. */
    interface TaskbarDesktopModeListener {
        /**
         * Callback for when task is resized in desktop mode.
         *
         * @param doesAnyTaskRequireTaskbarRounding whether task requires taskbar corner roundness.
         */
        fun onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding: Boolean) {}

        /**
         * Callback for when user is exiting desktop mode.
         *
         * @param duration for exit transition
         */
        fun onExitDesktopMode(duration: Int) {}

        /**
         * Callback for when user is entering desktop mode.
         *
         * @param duration for enter transition
         */
        fun onEnterDesktopMode(duration: Int) {}
    }

    companion object {
        @JvmField
        val INSTANCE = DaggerSingletonObject(LauncherAppComponent::getDesktopVisibilityController)

        private const val TAG = "DesktopVisController"
        private const val DEBUG = false

        const val INACTIVE_DESK_ID = -1
    }
}