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

import android.app.Activity
import android.app.Instrumentation
import android.content.ClipData
import android.content.Context
import android.content.Intent
import android.multiuser.Flags
import android.net.Uri
import android.os.statsd.privatespace.AddButtonShortcutClick
import android.platform.test.annotations.RequiresFlagsDisabled
import android.platform.test.annotations.RequiresFlagsEnabled
import android.platform.test.flag.junit.CheckFlagsRule
import android.platform.test.flag.junit.DeviceFlagsValueProvider
import android.util.StatsEvent
import android.util.StatsEventTestUtils
import android.util.StatsLog
import androidx.activity.result.launch
import androidx.lifecycle.Lifecycle
import androidx.test.core.app.ActivityScenario
import androidx.test.core.app.ApplicationProvider
import androidx.test.espresso.intent.Intents
import androidx.test.espresso.intent.matcher.IntentMatchers.hasAction
import androidx.test.espresso.intent.matcher.IntentMatchers.hasCategories
import androidx.test.espresso.intent.matcher.IntentMatchers.hasExtra
import androidx.test.espresso.intent.matcher.IntentMatchers.hasType
import androidx.test.espresso.intent.rule.IntentsRule
import com.android.bedstead.harrier.BedsteadJUnit4
import com.android.bedstead.multiuser.annotations.RequireRunOnPrivateProfile
import com.android.dx.mockito.inline.extended.ExtendedMockito
import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
import com.google.common.truth.Truth.assertThat
import com.google.protobuf.ExtensionRegistryLite
import java.util.stream.Collectors.toList
import kotlinx.coroutines.ExperimentalCoroutinesApi
import org.hamcrest.core.AllOf.allOf
import org.junit.After
import org.junit.Assert.assertThrows
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith
import org.mockito.Captor
import org.mockito.MockitoSession
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.atLeastOnce
import org.mockito.kotlin.mock
import org.mockito.kotlin.times
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever

@RequireRunOnPrivateProfile
@OptIn(ExperimentalCoroutinesApi::class)
@RunWith(BedsteadJUnit4::class)
class PrivateSpaceActivityTest {

    @get:Rule val intentsTestRule = IntentsRule()

    @get:Rule val checkFlagsRule: CheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule()

    @Captor val statsEventCaptor = argumentCaptor<StatsEvent>()

    private val registry = ExtensionRegistryLite.newInstance()

    private lateinit var mockitoSession: MockitoSession
    private lateinit var context: Context

    @Before
    fun setUp() {
        context = ApplicationProvider.getApplicationContext()
        mockitoSession = mockitoSession().mockStatic(StatsLog::class.java).startMocking()
        registry.add(PrivatespaceExtensionAtoms.privateSpaceAddButtonEvent)
    }

    @After
    fun tearDown() {
        mockitoSession.finishMocking()
    }

    @Test
    @RequiresFlagsDisabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionMain_movingContentFlagDisabled_throws() {
        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply { action = Intent.ACTION_MAIN }
        assertThrows(RuntimeException::class.java) {
            ActivityScenario.launch<PrivateSpaceActivity>(intent)
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionAddFiles_triggersDocumentPickerIntent() {
        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }
        ActivityScenario.launch<PrivateSpaceActivity>(intent).use { scenario ->
            Intents.intended(
                allOf(
                    hasAction(Intent.ACTION_OPEN_DOCUMENT),
                    hasCategories(setOf(Intent.CATEGORY_OPENABLE)),
                    hasType("*/*"),
                    hasExtra(Intent.EXTRA_ALLOW_MULTIPLE, true),
                    hasExtra(Intent.EXTRA_LOCAL_ONLY, true),
                )
            )
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionAddFiles_logsEvent() {

        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }
        ActivityScenario.launch<PrivateSpaceActivity>(intent)
        assertThat(statsEventCaptor).isNotNull()

        ExtendedMockito.verify({ StatsLog.write(statsEventCaptor.capture()) }, atLeastOnce())

        val atomList =
            statsEventCaptor.allValues
                .stream()
                .map { statsEvent -> StatsEventTestUtils.convertToAtom(statsEvent, registry) }
                .filter { e ->
                    e.hasExtension(PrivatespaceExtensionAtoms.privateSpaceAddButtonEvent)
                }
                .collect(toList())
        assertThat(atomList).hasSize(1)

        val atom = atomList[0]

        val privateSpaceAddButtonEvent =
            atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceAddButtonEvent)
        assertThat(privateSpaceAddButtonEvent.getAddButtonClick())
            .isEqualTo(AddButtonShortcutClick.ADD_BUTTON_CLICK_MOVE_FILES_SHORTCUT)
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionOpenMarketApp_finishesActivity() {
        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_OPEN_MARKET_APP
            }
        ActivityScenario.launch<PrivateSpaceActivity>(intent).use { scenario ->
            waitUntilActivityDestroyed(scenario)

            assertThat(scenario.state).isEqualTo(Lifecycle.State.DESTROYED)
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionOpenMarketApp_logsEvent() {

        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_OPEN_MARKET_APP
            }
        ActivityScenario.launch<PrivateSpaceActivity>(intent)
        assertThat(statsEventCaptor).isNotNull()

        ExtendedMockito.verify({ StatsLog.write(statsEventCaptor.capture()) }, atLeastOnce())

        val atomList =
            statsEventCaptor.allValues
                .stream()
                .map { statsEvent -> StatsEventTestUtils.convertToAtom(statsEvent, registry) }
                .filter { e ->
                    e.hasExtension(PrivatespaceExtensionAtoms.privateSpaceAddButtonEvent)
                }
                .collect(toList())
        assertThat(atomList).hasSize(1)

        val atom = atomList[0]

        val privateSpaceAddButtonEvent =
            atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceAddButtonEvent)
        assertThat(privateSpaceAddButtonEvent.getAddButtonClick())
            .isEqualTo(AddButtonShortcutClick.ADD_BUTTON_CLICK_OPEN_MARKET_APP_SHORTCUT)
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionUnknown_finishesActivity() {
        val intent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = "com.example.UNKNOWN_ACTION"
            }

        ActivityScenario.launch<PrivateSpaceActivity>(intent).use { scenario ->
            waitUntilActivityDestroyed(scenario)

            assertThat(scenario.state).isEqualTo(Lifecycle.State.DESTROYED)
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionNull_finishesActivity() {
        val intent = Intent(context, PrivateSpaceActivity::class.java) // action is null by default

        ActivityScenario.launch<PrivateSpaceActivity>(intent).use { scenario ->
            waitUntilActivityDestroyed(scenario)

            assertThat(scenario.state).isEqualTo(Lifecycle.State.DESTROYED)
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun intentActionAddFiles_activityResultNotOk_finishesActivity() {
        val launchIntent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }
        val resultData = Instrumentation.ActivityResult(Activity.RESULT_CANCELED, null)
        Intents.intending(hasAction(Intent.ACTION_OPEN_DOCUMENT)).respondWith(resultData)

        ActivityScenario.launch<PrivateSpaceActivity>(launchIntent).use { scenario ->
            assertThat(scenario.state).isEqualTo(Lifecycle.State.DESTROYED)
        }
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun documentsPickerResult_singleFile_updatesViewModelState() {
        val testUri: Uri = mock()
        whenever(testUri.userInfo).thenReturn("not null")
        val resultIntent = Intent().apply { data = testUri }
        val resultData = Instrumentation.ActivityResult(Activity.RESULT_OK, resultIntent)
        Intents.intending(hasAction(Intent.ACTION_OPEN_DOCUMENT)).respondWith(resultData)
        val launchIntent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }

        val scenario = ActivityScenario.launch<PrivateSpaceActivity>(launchIntent)

        scenario.onActivity { activity ->
            val viewModel = (activity as PrivateSpaceActivity).viewModel
            val uiState = viewModel.uiState
            assertThat(uiState).isEqualTo(PrivateSpaceUiState.SHOW_MOVE_FILES_DIALOG)
            assertThat(viewModel.fileUris).containsExactly(testUri)
        }

        scenario.close()
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun documentsPickerResult_multipleFiles_updatesViewModelState() {
        val testUri1: Uri = mock()
        val testUri2: Uri = mock()
        whenever(testUri1.userInfo).thenReturn("not null")
        whenever(testUri2.userInfo).thenReturn("not null")

        val resultIntent =
            Intent().apply {
                clipData =
                    ClipData.newUri(context.contentResolver, "files", testUri1).apply {
                        addItem(ClipData.Item(testUri2))
                    }
            }
        val resultData = Instrumentation.ActivityResult(Activity.RESULT_OK, resultIntent)

        Intents.intending(hasAction(Intent.ACTION_OPEN_DOCUMENT)).respondWith(resultData)

        val launchIntent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }
        val scenario = ActivityScenario.launch<PrivateSpaceActivity>(launchIntent)

        scenario.onActivity { activity ->
            val viewModel = (activity as PrivateSpaceActivity).viewModel
            val uiState = viewModel.uiState
            assertThat(uiState).isEqualTo(PrivateSpaceUiState.SHOW_MOVE_FILES_DIALOG)
            assertThat(viewModel.fileUris).containsExactly(testUri1, testUri2)
        }
        scenario.close()

        // assert only one intent to documents picker was launched
        assertThat(
                Intents.getIntents()
                    .stream()
                    .filter { intent -> intent.action == Intent.ACTION_OPEN_DOCUMENT }
                    .toList()
            )
            .hasSize(1)
    }

    @Test
    @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
    fun addFiles_recreateActivity_retainsDialogStateAndDoesNotRelaunchPicker() {
        val testUri1: Uri = mock()
        val testUri2: Uri = mock()
        whenever(testUri1.userInfo).thenReturn("not null")
        whenever(testUri2.userInfo).thenReturn("not null")
        val resultIntent =
            Intent().apply {
                clipData =
                    ClipData.newUri(context.contentResolver, "files", testUri1).apply {
                        addItem(ClipData.Item(testUri2))
                    }
            }
        val resultData = Instrumentation.ActivityResult(Activity.RESULT_OK, resultIntent)
        Intents.intending(hasAction(Intent.ACTION_OPEN_DOCUMENT)).respondWith(resultData)

        val launchIntent =
            Intent(context, PrivateSpaceActivity::class.java).apply {
                action = PrivateSpaceActivity.ACTION_ADD_FILES
            }
        val scenario = ActivityScenario.launch<PrivateSpaceActivity>(launchIntent)
        scenario.onActivity { activity ->
            val viewModel = (activity as PrivateSpaceActivity).viewModel
            val uiState = viewModel.uiState
            assertThat(uiState).isEqualTo(PrivateSpaceUiState.SHOW_MOVE_FILES_DIALOG)
            assertThat(viewModel.fileUris).containsExactly(testUri1, testUri2)
        }

        // Recreate the activity to simulate screen rotation.
        scenario.recreate()

        // Verify that the document picker was only launched once
        assertThat(
                Intents.getIntents()
                    .stream()
                    .filter { intent -> intent.action == Intent.ACTION_OPEN_DOCUMENT }
                    .toList()
            )
            .hasSize(1)

        scenario.close()
    }

    /** Helper function to wait for an ActivityScenario to reach the DESTROYED Lifecycle.State. */
    private fun <T : Activity> waitUntilActivityDestroyed(scenario: ActivityScenario<T>) {
        val pollIntervalMillis = 200L
        var pollRepetitions = 10 // total wait 2 seconds

        while (scenario.state != Lifecycle.State.DESTROYED) {
            Thread.sleep(pollIntervalMillis)
            if (--pollRepetitions <= 0)
                throw AssertionError(
                    "Activity Scenario did not reach state destroyed within expected time. " +
                        "Current state is '${scenario.state}'."
                )
        }
    }
}