package com.android.wallpaper.util.wallpaperconnection

import android.app.WallpaperInfo
import android.app.WallpaperManager
import android.app.wallpaper.WallpaperDescription
import android.content.ComponentName
import android.content.ContentValues
import android.content.Context
import android.content.Intent
import android.content.IntentFilter
import android.content.ServiceConnection
import android.graphics.Matrix
import android.graphics.Point
import android.net.Uri
import android.os.IBinder
import android.os.RemoteException
import android.service.wallpaper.IWallpaperEngine
import android.service.wallpaper.IWallpaperService
import android.service.wallpaper.WallpaperService
import android.util.Log
import android.view.MotionEvent
import android.view.SurfaceControl
import android.view.SurfaceView
import android.view.View
import com.android.app.tracing.TraceUtils.traceAsync
import com.android.wallpaper.effects.EffectsController
import com.android.wallpaper.model.Screen
import com.android.wallpaper.model.wallpaper.DeviceDisplayType
import com.android.wallpaper.picker.broadcast.BroadcastDispatcher
import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination.Companion.toSetWallpaperFlags
import com.android.wallpaper.picker.data.WallpaperModel.LiveWallpaperModel
import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
import com.android.wallpaper.util.ExtendedWallpaperEffectsUtils.isExtendedEffectWallpaper
import com.android.wallpaper.util.WallpaperConnection.WhichPreview
import dagger.hilt.android.qualifiers.ApplicationContext
import dagger.hilt.android.scopes.ActivityRetainedScoped
import java.lang.ref.WeakReference
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicBoolean
import javax.inject.Inject
import kotlinx.coroutines.CancellableContinuation
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.asExecutor
import kotlinx.coroutines.async
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

@ActivityRetainedScoped
class WallpaperConnectionUtils
@Inject
constructor(
    private val effectsController: EffectsController,
    @ApplicationContext private val context: Context,
    private val broadcastDispatcher: BroadcastDispatcher,
    @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
    @BackgroundDispatcher private val bgScope: CoroutineScope,
) {

    // The engineMap and the surfaceControlMap are used for disconnecting wallpaper services.
    private val wallpaperConnectionMap = ConcurrentHashMap<String, Deferred<WallpaperConnection>>()
    // Stores the latest connection for a service for later use like calling Engine methods.
    private val latestConnectionMap = ConcurrentHashMap<String, Deferred<WallpaperConnection>>()
    // Note that when one wallpaper engine's render is mirrored to a new surface view, we call
    // engine.mirrorSurfaceControl() and will have a new surface control instance.
    private val surfaceControlMap = mutableMapOf<String, MutableList<SurfaceControl>>()
    // Track the currently used creative wallpaper config preview URI to avoid unnecessary multiple
    // update queries for the same preview.
    private val creativeWallpaperConfigPreviewUriMap = mutableMapOf<String, Uri>()
    private val isPreviewEnginesConnected = CompletableDeferred<Boolean>()

    private val mutex = Mutex()

    private var disconnectOnWallpaperChange = false

    // Set true to enable debug logging (engine create/destroy, etc.)
    private val debug = false

    init {
        // We need to clear existing previews when a wallpaper is set since it's possible they will
        // be stale. See b/414490885.
        val wallpaperChanged =
            broadcastDispatcher.broadcastFlow(IntentFilter(Intent.ACTION_WALLPAPER_CHANGED))
        bgScope.launch {
            wallpaperChanged.collect { if (disconnectOnWallpaperChange) disconnectAll() }
        }
    }

    /**
     * Only call this function when the surface view is attached.
     *
     * @param totalEngineNum numbers of engines that should be considered connected when all of them
     *   are connected via [isPreviewEnginesConnected].
     */
    suspend fun connect(
        context: Context,
        wallpaperModel: LiveWallpaperModel,
        whichPreview: WhichPreview,
        destinationFlag: Int,
        surfaceView: SurfaceView,
        engineRenderingConfig: EngineRenderingConfig,
        isFirstBindingDeferred: CompletableDeferred<Boolean>,
        disconnectOnWallpaperChange: Boolean,
        totalEngineNum: Int = 1,
        listener: WallpaperEngineConnection.WallpaperEngineConnectionListener? = null,
        signalConfigChange: Boolean = false,
        onPreviewReady: (() -> Unit)? = null,
    ) {
        this.disconnectOnWallpaperChange = disconnectOnWallpaperChange
        val wallpaperInfo = wallpaperModel.liveWallpaperData.systemWallpaperInfo
        val engineDisplaySize = engineRenderingConfig.getEngineDisplaySize()
        val engineKey =
            wallpaperInfo.engineKey(
                engineDisplaySize,
                wallpaperModel.liveWallpaperData.description,
                wallpaperModel.liveWallpaperData.systemWallpaperInfo.component,
                destinationFlag,
            )
        val shortKey = engineKey.split(":").let { it.slice(2 until it.size) }.joinToString(":")

        if (debug) Log.d(TAG, "${toHex(this)}: requesting connection with key $engineKey")

        traceAsync(TAG, "connect") {
            // Update the creative wallpaper uri before starting the service.
            // We call this regardless of liveWallpaperContentHandling() because it's possible that
            // the flag is true here but false in the code we're calling.
            wallpaperModel.creativeWallpaperData?.configPreviewUri?.let {
                val uriKey =
                    wallpaperInfo.engineKey(
                        description = wallpaperModel.liveWallpaperData.description,
                        component = wallpaperModel.liveWallpaperData.systemWallpaperInfo.component,
                        destinationFlag = destinationFlag,
                    )
                if (!creativeWallpaperConfigPreviewUriMap.containsKey(uriKey)) {
                    mutex.withLock {
                        if (!creativeWallpaperConfigPreviewUriMap.containsKey(uriKey)) {
                            // First time binding wallpaper should initialize wallpaper preview.
                            if (isFirstBindingDeferred.await()) {
                                context.contentResolver.update(it, ContentValues(), null)
                            }
                            creativeWallpaperConfigPreviewUriMap[uriKey] = it
                        }
                    }
                }
            }

            val isConfigChange = signalConfigChange && !isFirstBindingDeferred.await()

            if (debug) Log.d(TAG, "isConfigChange: $isConfigChange")

            if (!wallpaperConnectionMap.containsKey(engineKey)) {
                mutex.withLock {
                    if (!wallpaperConnectionMap.containsKey(engineKey)) {
                        coroutineScope {
                            wallpaperConnectionMap[engineKey] = async {
                                initEngine(
                                    context,
                                    wallpaperModel.getWallpaperServiceIntent(),
                                    engineDisplaySize,
                                    destinationFlag,
                                    whichPreview,
                                    surfaceView,
                                    listener,
                                    wallpaperModel.liveWallpaperData.description,
                                    isConfigChange,
                                    this@WallpaperConnectionUtils,
                                    shortKey,
                                )
                            }

                            if (wallpaperConnectionMap.size == totalEngineNum) {
                                isPreviewEnginesConnected.complete(true)
                            }
                        }
                    }
                }
            }

            val serviceKey =
                wallpaperInfo.serviceKey(
                    wallpaperModel.liveWallpaperData.description,
                    wallpaperModel.liveWallpaperData.systemWallpaperInfo.component,
                )
            latestConnectionMap[serviceKey] =
                wallpaperConnectionMap[engineKey] as Deferred<WallpaperConnection>

            wallpaperConnectionMap[engineKey]?.await()?.let { connection ->
                if (debug) Log.d(TAG, "${toHex(this)}: MIR: $connection")
                connection.engineConnection.get()?.engine?.let {
                    mirrorAndReparent(
                        engineKey,
                        it,
                        surfaceView,
                        engineRenderingConfig.getEngineDisplaySize(),
                        engineRenderingConfig.enforceSingleEngine,
                        onPreviewReady,
                    )
                }
            }
        }
    }

    suspend fun disconnectAll() {
        if (debug) Log.d(TAG, "${toHex(this)}: disconnectAll()")
        surfaceControlMap.keys.map { key ->
            mutex.withLock {
                surfaceControlMap[key]?.let { surfaceControls ->
                    surfaceControls.forEach { it.release() }
                    surfaceControls.clear()
                }
            }
        }
        surfaceControlMap.clear()
        disconnectAllServices()
    }

    suspend fun disconnect(packageName: String) {
        mutex.withLock {
            if (debug) Log.d(TAG, "${toHex(this)}: disconnect(packageName)")
            surfaceControlMap.apply {
                filterKeys { key -> key.startsWith(packageName) }
                    .keys
                    .forEach { engineKey ->
                        remove(engineKey)?.let { surfaceControls ->
                            surfaceControls.forEach { it.release() }
                            surfaceControls.clear()
                        }
                        wallpaperConnectionMap.remove(engineKey)?.await()?.disconnect(context)
                    }
            }
        }
    }

    suspend fun disconnect(connection: WallpaperConnection) {
        mutex.withLock {
            if (debug) Log.d(TAG, "${toHex(this)}: disconnect(connection)")
            wallpaperConnectionMap
                .filterValues { engine ->
                    engine.await().engineConnection.get().let {
                        it != null && it == connection.engineConnection.get()
                    }
                }
                .keys
                .forEach { key ->
                    wallpaperConnectionMap.remove(key)?.await()?.disconnect(context)
                    surfaceControlMap.remove(key)?.let { surfaceControls ->
                        surfaceControls.forEach { it.release() }
                        surfaceControls.clear()
                    }
                }
        }
    }

    suspend fun setEngineVisibility(packageName: String, screen: Screen, isVisible: Boolean) {
        if (isPreviewEnginesConnected.await()) {
            mutex.withLock {
                wallpaperConnectionMap
                    .filterKeys { key ->
                        key.startsWith(packageName) && key.contains(":${screen.toFlag()}:")
                    }
                    .values
                    .forEach {
                        try {
                            it.await().engineConnection.get()?.engine?.setVisibility(isVisible)
                        } catch (e: RemoteException) {
                            Log.w(TAG, "Error setting engine visibility", e)
                        }
                    }
            }
        }
    }

    /**
     * Disconnect all live wallpaper services without releasing and clear surface controls. This
     * function is called before binding static wallpapers. We have cases that user switch between
     * live and static wallpapers. When switching from live to static wallpapers, we need to
     * disconnect the live wallpaper services to have the static wallpapers show up. But we can not
     * clear the surface controls yet, because we will need them to render the live wallpapers again
     * when switching from static to live wallpapers again.
     */
    suspend fun disconnectAllServices() {
        if (debug) Log.d(TAG, "${toHex(this)}: disconnectAllServices()")
        wallpaperConnectionMap.keys.map { key ->
            mutex.withLock { wallpaperConnectionMap.remove(key)?.await()?.disconnect(context) }
        }

        creativeWallpaperConfigPreviewUriMap.clear()
        latestConnectionMap.clear()
    }

    suspend fun dispatchTouchEvent(
        wallpaperModel: LiveWallpaperModel,
        engineRenderingConfig: EngineRenderingConfig,
        destinationFlag: Int,
        event: MotionEvent,
    ) {
        val engine =
            wallpaperModel.liveWallpaperData.systemWallpaperInfo
                .engineKey(
                    engineRenderingConfig.getEngineDisplaySize(),
                    wallpaperModel.liveWallpaperData.description,
                    wallpaperModel.liveWallpaperData.systemWallpaperInfo.component,
                    destinationFlag,
                )
                .let { engineKey ->
                    wallpaperConnectionMap[engineKey]?.await()?.engineConnection?.get()?.engine
                }

        if (engine != null) {
            val action: Int = event.actionMasked
            val dup = MotionEvent.obtainNoHistory(event).also { it.setLocation(event.x, event.y) }
            val pointerIndex = event.actionIndex
            try {
                engine.dispatchPointer(dup)
                if (action == MotionEvent.ACTION_UP) {
                    engine.dispatchWallpaperCommand(
                        WallpaperManager.COMMAND_TAP,
                        event.x.toInt(),
                        event.y.toInt(),
                        0,
                        null,
                    )
                } else if (action == MotionEvent.ACTION_POINTER_UP) {
                    engine.dispatchWallpaperCommand(
                        WallpaperManager.COMMAND_SECONDARY_TAP,
                        event.getX(pointerIndex).toInt(),
                        event.getY(pointerIndex).toInt(),
                        0,
                        null,
                    )
                }
            } catch (e: RemoteException) {
                Log.e(TAG, "Remote exception of wallpaper connection", e)
            }
        }
    }

    // Calls IWallpaperEngine#apply(which). Throws NoSuchMethodException if that method is not
    // defined, null if the Engine is not available, otherwise the result (which could also be
    // null).
    suspend fun applyWallpaper(
        destination: WallpaperDestination,
        wallpaperModel: LiveWallpaperModel,
    ): WallpaperDescription? {
        val wallpaperInfo = wallpaperModel.liveWallpaperData.systemWallpaperInfo
        val serviceKey =
            wallpaperInfo.serviceKey(
                wallpaperModel.liveWallpaperData.description,
                wallpaperModel.liveWallpaperData.systemWallpaperInfo.component,
            )
        latestConnectionMap[serviceKey]?.await()?.engineConnection?.get()?.engine?.let {
            try {
                return it.javaClass
                    .getMethod("onApplyWallpaper", Int::class.javaPrimitiveType)
                    .invoke(it, destination.toSetWallpaperFlags()) as WallpaperDescription?
            } catch (e: RemoteException) {
                // We catch this explicitly because it means that the method is defined, but the
                // bound object is dead.
                Log.w(TAG, "Error calling onApplyWallpaper", e)
            }
        }
        return null
    }

    private fun LiveWallpaperModel.getWallpaperServiceIntent(): Intent {
        return liveWallpaperData.systemWallpaperInfo.let {
            Intent(WallpaperService.SERVICE_INTERFACE).setClassName(it.packageName, it.serviceName)
        }
    }

    private suspend fun initEngine(
        context: Context,
        wallpaperIntent: Intent,
        displayMetrics: Point,
        destinationFlag: Int,
        whichPreview: WhichPreview,
        surfaceView: SurfaceView,
        listener: WallpaperEngineConnection.WallpaperEngineConnectionListener?,
        description: WallpaperDescription,
        isConfigChange: Boolean,
        owner: WallpaperConnectionUtils? = null,
        shortKey: String = "",
    ): WallpaperConnection {
        // Bind service and get service connection and wallpaper service
        val (serviceConnection, wallpaperService) = bindWallpaperService(context, wallpaperIntent)
        val engineConnection = WallpaperEngineConnection(displayMetrics, whichPreview)
        listener?.let { engineConnection.setListener(it) }
        val connection =
            WallpaperConnection(
                WeakReference(engineConnection),
                WeakReference(serviceConnection),
                WeakReference(wallpaperService),
                WeakReference(surfaceView.windowToken),
                debug,
            )
        (serviceConnection as WallpaperServiceConnection).deadConnectionListener =
            object : WallpaperServiceConnection.DeadConnectionListener {
                override fun onConnectionDead(serviceConnection: ServiceConnection) {
                    if (debug) Log.d(TAG, "${toHex(this)}: connection dead")
                    connection.disconnect(context)
                }
            }
        if (isConfigChange) description.content.apply { putBoolean(IS_CONFIG_CHANGE, true) }

        // Attach wallpaper connection to service and get wallpaper engine
        engineConnection
            .getEngine(wallpaperService, destinationFlag, surfaceView, description)
            .let { engine ->
                engine
                    .asBinder()
                    ?.linkToDeath(
                        {
                            if (debug) Log.d(TAG, "${toHex(this)}: engine dead")
                            connection.disconnect(context)
                        },
                        /* flags= */ 0,
                    )
            }
        surfaceView.viewTreeObserver.addOnWindowVisibilityChangeListener { visibility ->
            try {
                engineConnection.engine?.setVisibility(visibility == View.VISIBLE)
            } catch (e: RemoteException) {
                Log.w(TAG, "Error setting engine visibility", e)
            }
        }
        wallpaperService
            .asBinder()
            .linkToDeath(
                {
                    if (debug) Log.d(TAG, "${toHex(this)}: service dead")
                    connection.disconnect(context)
                },
                0,
            )

        if (debug) Log.d(TAG, "${toHex(owner)}: ADD: $connection: $shortKey")

        return connection
    }

    // Calculates a unique key for the wallpaper engine instance
    // TODO(b/390731022) Remove the MP-specific logic
    private fun WallpaperInfo.engineKey(
        displaySize: Point? = null,
        description: WallpaperDescription,
        component: ComponentName,
        destinationFlag: Int,
    ): String {
        // This is NOT the right way to do this long term. See b/390731022.
        val multiEngineExt =
            if (
                isExtendedEffectWallpaper(context, component) ||
                    isLegacyCinematicWallpaper(component)
            )
                ":$destinationFlag"
            else ""
        val keyWithoutSizeInformation =
            this.packageName
                .plus(":")
                .plus(this.serviceName)
                .plus(description.let { ":${it.id}" }.plus(multiEngineExt))
        return if (displaySize != null) {
            keyWithoutSizeInformation.plus(":").plus("${displaySize.x}x${displaySize.y}")
        } else {
            keyWithoutSizeInformation
        }
    }

    // Calculates a key unique to a service, but not a particular engine associated with that
    // service. Used as a key for latestConnectionMap.
    private fun WallpaperInfo.serviceKey(
        description: WallpaperDescription,
        component: ComponentName,
    ): String {
        return engineKey(null, description, component, 0)
    }

    private suspend fun bindWallpaperService(
        context: Context,
        intent: Intent,
    ): Pair<ServiceConnection, IWallpaperService> =
        suspendCancellableCoroutine {
            k: CancellableContinuation<Pair<ServiceConnection, IWallpaperService>> ->
            val serviceConnection =
                WallpaperServiceConnection(
                    object : WallpaperServiceConnection.WallpaperServiceConnectionListener {
                        override fun onWallpaperServiceConnected(
                            serviceConnection: ServiceConnection,
                            wallpaperService: IWallpaperService,
                        ) {
                            if (k.isActive) {
                                k.resumeWith(
                                    Result.success(Pair(serviceConnection, wallpaperService))
                                )
                            }
                        }
                    }
                )
            val success =
                context.bindService(
                    intent,
                    serviceConnection,
                    Context.BIND_AUTO_CREATE or
                        Context.BIND_IMPORTANT or
                        Context.BIND_ALLOW_ACTIVITY_STARTS,
                )
            if (!success && k.isActive) {
                k.resumeWith(Result.failure(Exception("Fail to bind the live wallpaper service.")))
            }
        }

    private suspend fun mirrorAndReparent(
        engineKey: String,
        engine: IWallpaperEngine,
        parentSurface: SurfaceView,
        displayMetrics: Point,
        enforceSingleEngine: Boolean,
        onPreviewReady: (() -> Unit)?,
    ) {
        fun logError(e: Exception) {
            Log.e(TAG, "Fail to reparent wallpaper surface", e)
        }

        try {
            val parentSurfaceControl = parentSurface.surfaceControl
            val wallpaperSurfaceControl = engine.mirrorSurfaceControl() ?: return
            // Add surface control reference for later release when disconnected
            addSurfaceControlReference(engineKey, wallpaperSurfaceControl)

            val values = getScale(parentSurface, displayMetrics)
            SurfaceControl.Transaction().use { t ->
                t.setMatrix(
                        wallpaperSurfaceControl,
                        if (enforceSingleEngine) values[Matrix.MSCALE_Y]
                        else values[Matrix.MSCALE_X],
                        values[Matrix.MSKEW_X],
                        values[Matrix.MSKEW_Y],
                        values[Matrix.MSCALE_Y],
                    )
                    .reparent(wallpaperSurfaceControl, parentSurfaceControl)
                    .show(wallpaperSurfaceControl)
                    .addTransactionCompletedListener(bgDispatcher.asExecutor()) { _ ->
                        onPreviewReady?.invoke()
                    }
                    .apply()
            }
        } catch (e: RemoteException) {
            logError(e)
        } catch (e: NullPointerException) {
            logError(e)
        }
    }

    private suspend fun addSurfaceControlReference(
        engineKey: String,
        wallpaperSurfaceControl: SurfaceControl,
    ) {
        val surfaceControls = surfaceControlMap[engineKey]
        if (surfaceControls == null) {
            mutex.withLock {
                surfaceControlMap[engineKey] =
                    (surfaceControlMap[engineKey] ?: mutableListOf()).apply {
                        add(wallpaperSurfaceControl)
                    }
            }
        } else {
            surfaceControls.add(wallpaperSurfaceControl)
        }
    }

    private fun getScale(parentSurface: SurfaceView, displayMetrics: Point): FloatArray {
        val metrics = Matrix()
        val values = FloatArray(9)
        val surfacePosition = parentSurface.holder.surfaceFrame
        metrics.postScale(
            surfacePosition.width().toFloat() / displayMetrics.x,
            surfacePosition.height().toFloat() / displayMetrics.y,
        )
        metrics.getValues(values)
        return values
    }

    private fun isLegacyCinematicWallpaper(component: ComponentName) =
        component.packageName == effectsController.effectsPackageName

    data class WallpaperConnection(
        val engineConnection: WeakReference<WallpaperEngineConnection>,
        val serviceConnection: WeakReference<ServiceConnection>,
        val wallpaperService: WeakReference<IWallpaperService>,
        val windowToken: WeakReference<IBinder>,
        val debug: Boolean = false,
    ) {
        private val disconnected = AtomicBoolean(false)

        fun disconnect(context: Context) {
            if (disconnected.compareAndSet(/* expectedValue= */ false, /* newValue= */ true)) {
                if (debug) Log.d(TAG, "DSC: $this")
                engineConnection.get()?.apply {
                    engine.let {
                        engine = null
                        try {
                            it?.destroy()
                        } catch (e: RemoteException) {
                            Log.w(TAG, "Error destroying wallpaper engine", e)
                        }
                        removeListener()
                    }
                }
                windowToken.get()?.let { window ->
                    wallpaperService.get()?.let { service ->
                        try {
                            service.detach(window)
                        } catch (e: RemoteException) {
                            Log.w(TAG, "Error disconnecting wallpaper service", e)
                        }
                    }
                }
                serviceConnection.get()?.let {
                    try {
                        context.unbindService(it)
                    } catch (e: RemoteException) {
                        Log.w(TAG, "Error unbinding service connection", e)
                    }
                }
            }
        }

        override fun toString(): String {
            return listOf(
                    "engineConnection serviceConnection wallpaperService windowToken",
                    toHex(engineConnection.get()),
                    toHex(serviceConnection.get()),
                    toHex(wallpaperService.get()),
                    toHex(windowToken.get()),
                    (serviceConnection.get() as? WallpaperServiceConnection)
                        ?.componentName
                        ?.packageName,
                )
                .joinToString()
        }
    }

    companion object {
        private const val TAG = "WallpaperConnectionUtils"
        private const val IS_CONFIG_CHANGE = "_picker_isConfigChange"

        data class EngineRenderingConfig(
            val enforceSingleEngine: Boolean,
            val deviceDisplayType: DeviceDisplayType,
            val smallDisplaySize: Point,
            val wallpaperDisplaySize: Point,
        ) {
            fun getEngineDisplaySize(): Point {
                // If we need to enforce single engine, always return the larger screen's preview
                return if (enforceSingleEngine) {
                    return wallpaperDisplaySize
                } else {
                    getPreviewDisplaySize()
                }
            }

            private fun getPreviewDisplaySize(): Point {
                return when (deviceDisplayType) {
                    DeviceDisplayType.SINGLE -> wallpaperDisplaySize
                    DeviceDisplayType.FOLDED -> smallDisplaySize
                    DeviceDisplayType.UNFOLDED -> wallpaperDisplaySize
                }
            }
        }

        fun LiveWallpaperModel.shouldEnforceSingleEngine(): Boolean {
            return when {
                creativeWallpaperData != null -> false
                liveWallpaperData.isEffectWallpaper -> false
                else -> !liveWallpaperData.supportsMultipleEngines
            }
        }

        fun toHex(o: Any?): String {
            return o.hashCode().toString(16).padStart(7, '0')
        }
    }
}