/*
 * 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.systemui.communal.widgets

import android.appwidget.AppWidgetHost
import android.appwidget.AppWidgetManager
import android.appwidget.AppWidgetProviderInfo
import android.content.ComponentName
import android.content.pm.UserInfo
import android.os.Bundle
import android.os.DeadObjectException
import android.os.UserHandle
import androidx.test.ext.junit.runners.AndroidJUnit4
import androidx.test.filters.SmallTest
import com.android.systemui.SysuiTestCase
import com.android.systemui.communal.shared.model.fakeGlanceableHubMultiUserHelper
import com.android.systemui.kosmos.applicationCoroutineScope
import com.android.systemui.kosmos.collectLastValue
import com.android.systemui.kosmos.collectValues
import com.android.systemui.kosmos.runTest
import com.android.systemui.kosmos.useUnconfinedTestDispatcher
import com.android.systemui.log.logcatLogBuffer
import com.android.systemui.testKosmos
import com.android.systemui.user.data.model.SelectedUserModel
import com.android.systemui.user.data.model.SelectionStatus
import com.android.systemui.user.data.repository.fakeUserRepository
import com.android.systemui.user.domain.interactor.SelectedUserInteractor
import com.android.systemui.user.domain.interactor.selectedUserInteractor
import com.android.systemui.util.mockito.any
import com.android.systemui.util.mockito.withArgCaptor
import com.google.common.truth.Truth.assertThat
import java.util.Optional
import org.junit.Before
import org.junit.Test
import org.junit.runner.RunWith
import org.mockito.ArgumentMatchers.eq
import org.mockito.Mock
import org.mockito.Mockito.never
import org.mockito.Mockito.times
import org.mockito.Mockito.verify
import org.mockito.MockitoAnnotations
import org.mockito.kotlin.whenever

@SmallTest
@RunWith(AndroidJUnit4::class)
class CommunalWidgetHostTest : SysuiTestCase() {
    private val kosmos = testKosmos().useUnconfinedTestDispatcher()

    @Mock private lateinit var appWidgetManager: AppWidgetManager
    @Mock private lateinit var appWidgetHost: CommunalAppWidgetHost
    @Mock private lateinit var providerInfo1: AppWidgetProviderInfo
    @Mock private lateinit var providerInfo2: AppWidgetProviderInfo
    @Mock private lateinit var providerInfo3: AppWidgetProviderInfo

    private val selectedUserInteractor: SelectedUserInteractor by lazy {
        kosmos.selectedUserInteractor
    }

    private lateinit var underTest: CommunalWidgetHost

    @Before
    fun setUp() {
        MockitoAnnotations.initMocks(this)

        whenever(
                appWidgetManager.bindAppWidgetIdIfAllowed(
                    any<Int>(),
                    any<UserHandle>(),
                    any<ComponentName>(),
                    any<Bundle>(),
                )
            )
            .thenReturn(true)

        underTest =
            CommunalWidgetHost(
                kosmos.applicationCoroutineScope,
                Optional.of(appWidgetManager),
                appWidgetHost,
                selectedUserInteractor,
                logcatLogBuffer("CommunalWidgetHostTest"),
                glanceableHubMultiUserHelper = kosmos.fakeGlanceableHubMultiUserHelper,
            )
    }

    @Test
    fun allocateIdAndBindWidget_withCurrentUser() =
        kosmos.runTest {
            val provider = ComponentName("pkg_name", "cls_name")
            val widgetId = 1
            val userId by collectLastValue(selectedUserInteractor.selectedUser)
            selectUser()

            val user = UserHandle(checkNotNull(userId))
            whenever(appWidgetHost.allocateAppWidgetId()).thenReturn(widgetId)

            // bind the widget with the current user when no user is explicitly set
            val result = underTest.allocateIdAndBindWidget(provider)

            verify(appWidgetHost).allocateAppWidgetId()
            val bundle =
                withArgCaptor<Bundle> {
                    verify(appWidgetManager)
                        .bindAppWidgetIdIfAllowed(eq(widgetId), eq(user), eq(provider), capture())
                }
            assertThat(result).isEqualTo(widgetId)
            assertThat(bundle.getInt(AppWidgetManager.OPTION_APPWIDGET_HOST_CATEGORY))
                .isEqualTo(AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD)
        }

    @Test
    fun allocateIdAndBindWidget_onSuccess() =
        kosmos.runTest {
            val provider = ComponentName("pkg_name", "cls_name")
            val widgetId = 1
            val user = UserHandle(0)

            whenever(appWidgetHost.allocateAppWidgetId()).thenReturn(widgetId)

            // provider and user handle are both set
            val result = underTest.allocateIdAndBindWidget(provider, user)

            verify(appWidgetHost).allocateAppWidgetId()
            val bundle =
                withArgCaptor<Bundle> {
                    verify(appWidgetManager)
                        .bindAppWidgetIdIfAllowed(eq(widgetId), eq(user), eq(provider), capture())
                }
            assertThat(result).isEqualTo(widgetId)
            assertThat(bundle.getInt(AppWidgetManager.OPTION_APPWIDGET_HOST_CATEGORY))
                .isEqualTo(AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD)
        }

    @Test
    fun allocateIdAndBindWidget_onFailure() =
        kosmos.runTest {
            val provider = ComponentName("pkg_name", "cls_name")
            val widgetId = 1
            val user = UserHandle(0)

            whenever(appWidgetHost.allocateAppWidgetId()).thenReturn(widgetId)
            // failed to bind widget
            whenever(
                    appWidgetManager.bindAppWidgetIdIfAllowed(
                        any<Int>(),
                        any<UserHandle>(),
                        any<ComponentName>(),
                        any<Bundle>(),
                    )
                )
                .thenReturn(false)
            val result = underTest.allocateIdAndBindWidget(provider, user)

            verify(appWidgetHost).allocateAppWidgetId()
            verify(appWidgetManager)
                .bindAppWidgetIdIfAllowed(eq(widgetId), eq(user), eq(provider), any())
            verify(appWidgetHost).deleteAppWidgetId(widgetId)
            assertThat(result).isNull()
        }

    @Test
    fun listener_exactlyOneListenerRegisteredForEachWidgetWhenHostStartListening() =
        kosmos.runTest {
            // 3 widgets registered with the host
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2, 3))

            underTest.startObservingHost()

            // Make sure no listener is set before host starts listening
            verify(appWidgetHost, never()).setListener(any(), any())

            // Host starts listening
            val observer =
                withArgCaptor<CommunalAppWidgetHost.Observer> {
                    verify(appWidgetHost).addObserver(capture())
                }
            observer.onHostStartListening()

            // Verify a listener is set for each widget
            verify(appWidgetHost, times(3)).setListener(any(), any())
            verify(appWidgetHost).setListener(eq(1), any())
            verify(appWidgetHost).setListener(eq(2), any())
            verify(appWidgetHost).setListener(eq(3), any())
        }

    @Test
    fun listener_listenersRemovedWhenHostStopListening() =
        kosmos.runTest {
            // 3 widgets registered with the host
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2, 3))

            underTest.startObservingHost()

            // Host starts listening
            val observer =
                withArgCaptor<CommunalAppWidgetHost.Observer> {
                    verify(appWidgetHost).addObserver(capture())
                }
            observer.onHostStartListening()

            // Verify none of the listener is removed before host stop listening
            verify(appWidgetHost, never()).removeListener(any())

            observer.onHostStopListening()

            // Verify each listener is removed
            verify(appWidgetHost, times(3)).removeListener(any())
            verify(appWidgetHost).removeListener(eq(1))
            verify(appWidgetHost).removeListener(eq(2))
            verify(appWidgetHost).removeListener(eq(3))
        }

    @Test
    fun listener_addNewListenerWhenNewIdAllocated() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf())
            val observer = start()

            // Verify no listener is set before a new app widget id is allocated
            verify(appWidgetHost, never()).setListener(any(), any())

            // Allocate an app widget id
            observer.onAllocateAppWidgetId(1)

            // Verify new listener set for that app widget id
            verify(appWidgetHost).setListener(eq(1), any())
        }

    @Test
    fun listener_removeListenerWhenWidgetDeleted() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1))
            val observer = start()

            // Verify listener not removed before widget deleted
            verify(appWidgetHost, never()).removeListener(eq(1))

            // Widget deleted
            observer.onDeleteAppWidgetId(1)

            // Verify listener removed for that widget
            verify(appWidgetHost).removeListener(eq(1))
        }

    @Test
    fun providerInfo_populatesWhenStartListening() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)

            val providerInfoValues by collectValues(underTest.appWidgetProviders)

            // Assert that the map is empty before host starts listening
            assertThat(providerInfoValues).hasSize(1)
            assertThat(providerInfoValues[0]).isEmpty()

            start()

            // Assert that the provider info map is populated after host started listening, and that
            // all providers are emitted at once
            assertThat(providerInfoValues).hasSize(2)
            assertThat(providerInfoValues[1])
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2)))
        }

    @Test
    fun providerInfo_clearsWhenStopListening() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)

            val observer = start()

            // Assert that the provider info map is populated
            val providerInfo by collectLastValue(underTest.appWidgetProviders)
            assertThat(providerInfo)
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2)))

            // Host stop listening
            observer.onHostStopListening()

            // Assert that the provider info map is cleared
            assertThat(providerInfo).isEmpty()
        }

    @Test
    fun providerInfo_onUpdate() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)

            val providerInfo by collectLastValue(underTest.appWidgetProviders)

            start()

            // Assert that the provider info map is populated
            assertThat(providerInfo)
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2)))

            // Provider info for widget 1 updated
            val listener =
                withArgCaptor<AppWidgetHost.AppWidgetHostListener> {
                    verify(appWidgetHost).setListener(eq(1), capture())
                }
            listener.onUpdateProviderInfo(providerInfo3)

            // Assert that the update is reflected in the flow
            assertThat(providerInfo)
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo3), Pair(2, providerInfo2)))
        }

    @Test
    fun providerInfo_updateWhenANewWidgetIsBound() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)

            val providerInfo by collectLastValue(underTest.appWidgetProviders)

            start()

            // Assert that the provider info map is populated
            assertThat(providerInfo)
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2)))

            // Bind a new widget
            whenever(appWidgetHost.allocateAppWidgetId()).thenReturn(3)
            whenever(appWidgetManager.getAppWidgetInfo(3)).thenReturn(providerInfo3)
            val newWidgetComponentName = ComponentName.unflattenFromString("pkg_new/cls_new")!!
            underTest.allocateIdAndBindWidget(newWidgetComponentName)

            // Assert that the new provider is reflected in the flow
            assertThat(providerInfo)
                .containsExactlyEntriesIn(
                    mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2), Pair(3, providerInfo3))
                )
        }

    @Test
    fun providerInfo_updateWhenWidgetRemoved() =
        kosmos.runTest {
            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)

            val providerInfo by collectLastValue(underTest.appWidgetProviders)

            val observer = start()

            // Assert that the provider info map is populated
            assertThat(providerInfo)
                .containsExactlyEntriesIn(mapOf(Pair(1, providerInfo1), Pair(2, providerInfo2)))

            // Remove widget 1
            observer.onDeleteAppWidgetId(1)

            // Assert that provider info for widget 1 is removed
            assertThat(providerInfo).containsExactlyEntriesIn(mapOf(Pair(2, providerInfo2)))
        }

    @Test
    fun refreshProviderInfo_handlesDeadObjectException() =
        kosmos.runTest {
            val providerInfoValues by collectLastValue(underTest.appWidgetProviders)

            whenever(appWidgetHost.appWidgetIds).thenReturn(intArrayOf(1, 2))
            whenever(appWidgetHost.setListener(eq(1), org.mockito.kotlin.any())).thenAnswer {
                throw DeadObjectException()
            }
            whenever(appWidgetManager.getAppWidgetInfo(1)).thenReturn(providerInfo1)
            whenever(appWidgetManager.getAppWidgetInfo(2)).thenReturn(providerInfo2)
            underTest.refreshProviders()

            with(providerInfoValues!!) {
                assertThat(size).isEqualTo(1)
                assertThat(containsKey(2)).isTrue()
            }
        }

    @Test
    fun onAllocateAppWidgetId_handlesDeadObjectException() =
        kosmos.runTest {
            whenever(appWidgetHost.setListener(eq(1), org.mockito.kotlin.any())).thenAnswer {
                throw DeadObjectException()
            }
            underTest.onAllocateAppWidgetId(1)
        }

    private fun selectUser() {
        kosmos.fakeUserRepository.selectedUser.value =
            SelectedUserModel(
                userInfo = UserInfo(0, "Current user", 0),
                selectionStatus = SelectionStatus.SELECTION_COMPLETE,
            )
    }

    private fun start(): CommunalAppWidgetHost.Observer {
        underTest.startObservingHost()

        val observer =
            withArgCaptor<CommunalAppWidgetHost.Observer> {
                verify(appWidgetHost).addObserver(capture())
            }
        observer.onHostStartListening()
        return observer
    }
}