/*
 * 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.desktop

import android.content.Intent
import android.content.res.Configuration
import android.graphics.Bitmap
import android.hardware.display.DisplayManager
import android.os.Binder
import android.os.IBinder
import android.util.Log
import android.view.WindowManager
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.platform.ComposeView
import androidx.lifecycle.LifecycleService
import androidx.lifecycle.SAVED_STATE_REGISTRY_OWNER_KEY
import androidx.lifecycle.VIEW_MODEL_STORE_OWNER_KEY
import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.ViewModelStore
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.createSavedStateHandle
import androidx.lifecycle.enableSavedStateHandles
import androidx.lifecycle.lifecycleScope
import androidx.lifecycle.setViewTreeLifecycleOwner
import androidx.lifecycle.setViewTreeViewModelStoreOwner
import androidx.lifecycle.viewmodel.CreationExtras
import androidx.lifecycle.viewmodel.MutableCreationExtras
import androidx.savedstate.SavedStateRegistry
import androidx.savedstate.SavedStateRegistryController
import androidx.savedstate.SavedStateRegistryOwner
import androidx.savedstate.setViewTreeSavedStateRegistryOwner
import com.android.eyedropper.EyeDropperController
import com.android.eyedropper.ui.desktop.DesktopSelectColorOverlay
import com.android.eyedropper.ui.desktop.DesktopViewModel
import com.android.eyedropper.util.InputHelper
import com.android.eyedropper.util.WindowHelper
import dagger.hilt.android.AndroidEntryPoint
import javax.inject.Inject
import kotlin.reflect.KClass
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.launch
private const val TAG = "DesktopController" // DesktopControllerService too long

// Bundle UI resources needed per display
private data class DisplayComponents(
    val windowManager: WindowManager,
    val selectColorView: ComposeView,
    val viewModel: DesktopViewModel
)

/**
 * Service bound to by [MainActivity] responsible for displaying and managing the Desktop UI.
 *
 * Extends `LifecycleService` because the Compose runtime must attach to a proper system
 * managed LifecycleOwner. Unlike Activities/Fragments, we must manually do this for a service.
 *
 * This service implements multi-display. This is achieved by managing per-display WindowContext,
 * WindowManager, and DesktopViewModel resources. No display knows about any other display.
 */
@AndroidEntryPoint(LifecycleService::class)
class DesktopControllerService : Hilt_DesktopControllerService(),
    EyeDropperController,
    ViewModelStoreOwner,
    SavedStateRegistryOwner
{
    @Inject lateinit var inputHelper: InputHelper
    @Inject lateinit var windowHelper: WindowHelper
    @Inject lateinit var inputDeviceListener: DesktopInputDeviceListener
    @Inject lateinit var displayListener: DesktopDisplayListener
    @Inject lateinit var displayManager: DisplayManager

    private var abortMainActivityCallback: (() -> Unit)? = null // needed for onConfigChanged
    private var hasOverlayShown = false // "Overlay" refers to all overlays across all displays

    private lateinit var displayIds: Set<Int>
    lateinit var displayIdToDeferredScreenshotMap: Map<Int, Deferred<Bitmap>>
    private val displayIdToDisplayComponentsMap: MutableMap<Int, DisplayComponents> = mutableMapOf()

    /**
     * ViewModelStoreOwner/SavedStateRegistryOwner implementations required for LifecycleService.
     * These are unused in multi-display implementation, since each display is mapped to its own
     * ViewModelStoreOwner instance.
     */
    private val _viewModelStore = ViewModelStore()
    override val viewModelStore: ViewModelStore
        get() = _viewModelStore

    private val savedStateRegistryController = SavedStateRegistryController.create(this)
    override val savedStateRegistry: SavedStateRegistry
        get() = savedStateRegistryController.savedStateRegistry

    /**
     * Custom ViewModelFactory to construct [DesktopViewModel] with SavedStateHandle.
     * Necessary because there is no built-in method to pull a ViewModel into a Service.
     */
    private inner class DesktopViewModelFactory : ViewModelProvider.Factory {
        override fun <T : ViewModel> create(modelClass: KClass<T>, extras: CreationExtras): T {
            val savedStateHandle = extras.createSavedStateHandle()
            @Suppress("UNCHECKED_CAST")
            return DesktopViewModel(savedStateHandle) as T
        }
    }

    // LocalBinder to communicate with Activity client, which will always run in the same process.
    private val binder = LocalBinder()
    inner class LocalBinder : Binder() {
        fun getService(): DesktopControllerService = this@DesktopControllerService
    }

    override fun onBind(intent: Intent): IBinder {
        Log.d(TAG, "onBind called")
        super.onBind(intent)
        return binder
    }

    override fun onCreate() {
        Log.d(TAG, "onCreate called")
        super.onCreate()
        enableSavedStateHandles()

        savedStateRegistryController.run {
            performAttach() // Connect registry to the lifecycle
            performRestore(null)
        }
    }

    // TODO: b/438271594 - Make configuration changes launch fresh screenshot
    // This is a temporary remediation that just aborts the request on configuration change.
    override fun onConfigurationChanged(newConfig: Configuration) {
        super.onConfigurationChanged(newConfig)
        Log.d(TAG, "onConfigurationChanged called, newConfig: $newConfig")
        abortMainActivityCallback?.invoke()
    }

    override fun onDestroy() {
        Log.d(TAG, "onDestroy called")
        super.onDestroy()
        removeUiOverlay()
    }

    /**
     * For responsiveness, force callbacks to immediately remove reticle on select/abort. Flow is:
     * <make selection> -> removeUiOverlay -> MainActivity: sendResult -> finish -> onStop/Destroy
     *
     * Multi-display support unfortunately makes this function more expensive. This is because
     * multi-display resource instantiation cannot start until we know the `displayIds`, which is
     * only set through the LocalBinder after onCreate() is already called. This service shouldn't
     * obtain its own set of `displayIds` to ensure against inconsistent states.
     */
    override fun showUiOverlay(onSelectColor: (Int) -> Unit, onAbort: () -> Unit) {
        Log.d(TAG, "hasOverlayShown = $hasOverlayShown, showUiOverlay called")
        if (hasOverlayShown) return
        hasOverlayShown = true

        displayIds = displayIdToDeferredScreenshotMap.keys
        if (displayIds.isEmpty()) throw RuntimeException("Desktop mode finds no displays present")

        val wrappedOnSelectColor = { colorInt: Int ->
            removeUiOverlay()
            onSelectColor(colorInt)
        }
        val wrappedOnAbort = {
            removeUiOverlay()
            onAbort()
        }

        // While polling for the cursor here does make it possible for the user to have switched
        // displays between EyeDropper launch and UI display, it does ensure consistency
        // (guarantees that only one cursor rendered), which is not true if we poll any later, or
        // if each display tries to poll individually.
        val cursorPositionPair = inputHelper.getCursorPosition(displayIds)
            ?: throw RuntimeException("Desktop mode finds no cursors present.")

        val (displayWithCursorId, cursorPosition) = cursorPositionPair

        displayIds.forEach { displayId ->
            val cursorPositionOrNull = cursorPosition.takeIf { displayId == displayWithCursorId }
            setupDisplayComponents(displayId)
            setupViewResources(
                displayId,
                cursorPositionOrNull,
                wrappedOnSelectColor,
                wrappedOnAbort
            )
        }
        // Only one listener between all displays
        displayListener.initialize(wrappedOnAbort)
        inputDeviceListener.initialize(wrappedOnAbort)

        displayManager.registerDisplayListener(displayListener, null)
        inputHelper.registerInputDeviceListener(inputDeviceListener)
    }

    override fun removeUiOverlay() {
        Log.d(TAG, "hasOverlayShown = $hasOverlayShown, removeUIOverlay called")
        if (!hasOverlayShown) return
        hasOverlayShown = false

        displayIds.forEach { displayId -> teardownViewResources(displayId) }

        displayManager.unregisterDisplayListener(displayListener)
        inputHelper.unregisterInputDeviceListener(inputDeviceListener)
    }

    // Initialize `displayIdToDisplayComponentsMap`.
    private fun setupDisplayComponents(displayId: Int) {
        Log.d(TAG, "Initializing display components for display $displayId...")

        val windowContext = windowHelper.getWindowContext(this, displayId)
        val windowManager = windowHelper.getWindowManager(windowContext)

        val selectColorView = ComposeView(windowContext).apply {
            setViewTreeLifecycleOwner(this@DesktopControllerService)
            setViewTreeViewModelStoreOwner(this@DesktopControllerService)
            setViewTreeSavedStateRegistryOwner(this@DesktopControllerService)
        }

        val extras = MutableCreationExtras().apply {
            this[SAVED_STATE_REGISTRY_OWNER_KEY] = this@DesktopControllerService
            this[VIEW_MODEL_STORE_OWNER_KEY] = this@DesktopControllerService
            this[ViewModelProvider.AndroidViewModelFactory.APPLICATION_KEY] = application
        }
        val viewModel = ViewModelProvider(
            ViewModelStore(),
            DesktopViewModelFactory(),
            extras,
        )[DesktopViewModel::class.java]

        displayIdToDisplayComponentsMap[displayId] = DisplayComponents(
            windowManager, selectColorView, viewModel
        )
    }

    // Await screenshot just-in-time for viewModel and View initialization
    private fun setupViewResources(
        displayId: Int,
        cursorPosition: Offset?,
        onColorSelected: (Int) -> Unit,
        onSelectionCanceled: () -> Unit,
    ) {
        Log.d(TAG, "Initializing overlay view for display $displayId...")

        lifecycleScope.launch {
            val deferredScreenshotBitmap = displayIdToDeferredScreenshotMap[displayId]!!
            val displayComponent = displayIdToDisplayComponentsMap[displayId]!!

            val selectColorView = displayComponent.selectColorView
            val windowManager = displayComponent.windowManager

            val viewModel = displayComponent.viewModel
            viewModel.initialize(
                deferredScreenshotBitmap.await(),
                cursorPosition,
                onColorSelected,
                onSelectionCanceled
            )

            try {
                selectColorView.setContent {
                    DesktopSelectColorOverlay(viewModel)
                }
                windowHelper.addFocusedTransparentOverlay(windowManager, selectColorView)
            } catch (e: Exception) {
                Log.e(TAG, "Exception creating and adding overlay view to display $displayId: $e")
            }
        }
    }

    private fun teardownViewResources(displayId: Int) {
        Log.d(TAG, "Removing overlay for display $displayId...")
        val displayComponent = displayIdToDisplayComponentsMap[displayId]!!
        val windowManager = displayComponent.windowManager
        val selectColorView = displayComponent.selectColorView

        try {
            windowHelper.removeView(windowManager, selectColorView)
        } catch (e: Exception) {
            Log.e(TAG, "Error removing overlay on display $displayId: $e")
        }
    }
}