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

import android.content.ContentResolver
import android.content.Context
import android.content.ContextWrapper
import android.database.Cursor
import android.multiuser.Flags
import android.net.Uri
import android.os.statsd.privatespace.TransferErrorCode
import android.os.statsd.privatespace.TransferResult
import android.platform.test.annotations.RequiresFlagsEnabled
import android.platform.test.flag.junit.CheckFlagsRule
import android.platform.test.flag.junit.DeviceFlagsValueProvider
import android.provider.DocumentsContract
import android.util.StatsEvent
import android.util.StatsEventTestUtils
import android.util.StatsLog
import androidx.datastore.core.IOException
import androidx.test.platform.app.InstrumentationRegistry
import com.android.bedstead.harrier.BedsteadJUnit4
import com.android.dx.mockito.inline.extended.ExtendedMockito
import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
import com.android.privatespace.PrivateSpaceMoveContentEvent
import com.android.privatespace.PrivatespaceExtensionAtoms
import com.android.privatespace.StorageHelper
import com.google.common.truth.Truth.assertThat
import com.google.protobuf.ExtensionRegistryLite
import java.io.FileNotFoundException
import java.io.InputStream
import java.io.OutputStream
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.test.runTest
import org.junit.After
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith
import org.mockito.Mockito.inOrder
import org.mockito.Mockito.verify
import org.mockito.MockitoSession
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.eq
import org.mockito.kotlin.mock
import org.mockito.kotlin.never
import org.mockito.kotlin.whenever
import org.mockito.quality.Strictness

private const val PATH = "PATH"

@RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
@RunWith(BedsteadJUnit4::class)
class FileTransferManagerImplTest {
    @get:Rule val checkFlagsRule: CheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule()

    private lateinit var fileTransferManagerImpl: FileTransferManagerImpl

    private lateinit var mockContext: Context
    private var realContext: Context = InstrumentationRegistry.getInstrumentation().targetContext

    private val mockNotificationsHelper: NotificationsHelper = mock()
    private val mockContentResolver: ContentResolver = mock()
    private val mockFileTransferStateRepository: IFileTransferStateRepository = mock()
    private val mockStorageHelper: StorageHelper = mock()
    private val mockUri: Uri = mock()
    private val mockCursor: Cursor = mock()
    private val mockInsertUri: Uri = mock()
    private val mockInputStream: InputStream = mock()
    private val mockOutputStream: OutputStream = mock()
    private val mockClock: SystemClockWrapper = mock()
    private lateinit var mockitoSession: MockitoSession
    private val statsEventCaptor = argumentCaptor<StatsEvent>()
    private val registry = ExtensionRegistryLite.newInstance()

    private val uris = listOf(mockUri)

    @Before
    fun setUp() = runBlocking {
        mockitoSession =
            mockitoSession()
                .strictness(Strictness.LENIENT)
                .mockStatic(StatsLog::class.java)
                .spyStatic(DocumentsContract::class.java)
                .startMocking()
        registry.add(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)

        mockContext =
            object : ContextWrapper(realContext) {
                override fun getContentResolver(): ContentResolver? {
                    return mockContentResolver
                }
            }
        whenever(mockStorageHelper.calculateTotalSize(any(), any(), any())).thenReturn(1)
        whenever(mockStorageHelper.getAvailableDeviceStorage()).thenReturn(Long.MAX_VALUE)
        whenever(mockContentResolver.query(any(), any(), eq(null), eq(null))).thenReturn(mockCursor)
        whenever(mockCursor.moveToFirst()).thenReturn(true)
        whenever(mockCursor.getString(any())).thenReturn("test")
        whenever(mockContentResolver.insert(any(), any())).thenReturn(mockInsertUri)
        whenever(mockContentResolver.openOutputStream(any())).thenReturn(mockOutputStream)
        whenever(mockContentResolver.openInputStream(any())).thenReturn(mockInputStream)

        fileTransferManagerImpl =
            FileTransferManagerImpl(
                notificationsHelper = mockNotificationsHelper,
                context = mockContext,
                fileTransferStateRepository = mockFileTransferStateRepository,
                storageHelper = mockStorageHelper,
                clock = mockClock,
            )
    }

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

    @Test
    fun transferFiles_success_displaysSuccessfulCompletionNotification() = runTest {
        val numberOfFiles = 2
        val keepOriginal = true
        fileTransferManagerImpl.transferFiles(uris, keepOriginal, PATH)

        fileTransferManagerImpl.postEndOfTransferNotification(numberOfFiles, keepOriginal)

        verify(mockNotificationsHelper)
            .displaySuccessfulCompletionNotification(numberOfFiles, keepOriginal)
        verify(mockNotificationsHelper, never()).displayPartialTransferErrorNotification(any())
    }

    @Test
    fun transferFiles_error_displaysPartialTransferErrorNotification() = runTest {
        val numberOfFiles = 2
        val keepOriginal = true
        whenever(mockContentResolver.insert(any(), any())).thenReturn(null)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        fileTransferManagerImpl.postEndOfTransferNotification(numberOfFiles, keepOriginal)

        verify(mockNotificationsHelper).displayPartialTransferErrorNotification(keepOriginal)
        verify(mockNotificationsHelper, never())
            .displaySuccessfulCompletionNotification(any(), any())
    }

    @Test
    fun transferFiles_errorAndSuccess_displayPartialTransferErrorNotification() = runTest {
        val numberOfFiles = 2
        val keepOriginal = true
        val mockUris: List<Uri> = List(3) { mock() }
        ExtendedMockito.doThrow(FileNotFoundException()).`when` {
            DocumentsContract.deleteDocument(any(), eq(mockUris[0]))
        }
        fileTransferManagerImpl.transferFiles(mockUris, false, PATH)

        fileTransferManagerImpl.postEndOfTransferNotification(numberOfFiles, keepOriginal)

        verify(mockNotificationsHelper).displayPartialTransferErrorNotification(keepOriginal)
        verify(mockNotificationsHelper, never())
            .displaySuccessfulCompletionNotification(any(), any())
    }

    @Test
    fun transferFiles_success_logsEvent() = runTest {
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_SUCCESS)
        assertThat(event.successFileMimeTypesCount).isEqualTo(1)
    }

    @Test
    fun transferFiles_aboveAvailableDeviceStorage_logsEvent() = runTest {
        whenever(mockStorageHelper.getAvailableDeviceStorage()).thenReturn(0)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE)
    }

    @Test
    fun transferFiles_errorGetMetadata_logsEvent() = runTest {
        whenever(mockCursor.moveToFirst()).thenReturn(false)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_GET_METADATA)
    }

    @Test
    fun transferFiles_errorCreateNewMediaEntry_logsEvent() = runTest {
        whenever(mockContentResolver.insert(any(), any())).thenReturn(null)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_CREATE_NEW_MEDIA_ENTRY)
    }

    @Test
    fun transferFiles_errorGetFileName_logsEvent() = runTest {
        whenever(mockCursor.getString(any())).thenThrow(IllegalArgumentException())
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_GET_FILE_NAME)
    }

    @Test
    fun transferFiles_errorQueryUri_logsEvent() = runTest {
        whenever(mockContentResolver.query(any(), any(), eq(null), eq(null))).thenReturn(null)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_QUERY_URI)
    }

    @Test
    fun transferFiles_errorOpenOutputStream_logsEvent() = runTest {
        whenever(mockContentResolver.openOutputStream(any())).thenReturn(null)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)

        verify(mockContentResolver).openInputStream(any())
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_OPEN_OUTPUT_STREAM)
    }

    @Test
    fun transferFiles_errorOpenInputStream_logsEvent() = runTest {
        whenever(mockContentResolver.openInputStream(any())).thenReturn(null)
        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_OPEN_INPUT_STREAM)
    }

    @Test
    fun transferFiles_errorCopyFile_logsEvent() = runTest {
        whenever(mockInputStream.read(any())).thenThrow(IOException())

        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_COPY_FILE)
    }

    @Test
    fun transferFiles_errorRemoveOriginalFile_logsEvent() = runTest {
        ExtendedMockito.doThrow(FileNotFoundException()).`when` {
            DocumentsContract.deleteDocument(any(), any())
        }

        fileTransferManagerImpl.transferFiles(uris, false, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE)
    }

    @Test
    fun transferFiles_multipleErrorsAndSuccess_logsEvent() = runTest {
        val mockUris: List<Uri> = List(4) { mock() }

        // Expect TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE for mockUris[0]
        ExtendedMockito.doThrow(FileNotFoundException()).`when` {
            DocumentsContract.deleteDocument(any(), eq(mockUris[0]))
        }
        // Expect TransferErrorCode.TRANSFER_ERROR_QUERY_URI for mockUris[1]
        whenever(mockContentResolver.query(eq(mockUris[1]), any(), eq(null), eq(null)))
            .thenReturn(null)

        fileTransferManagerImpl.transferFiles(mockUris, false, PATH)

        val event = getPrivateSpaceMoveContentStatsEvent()
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_PARTIAL_SUCCESS)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(
                TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE,
                TransferErrorCode.TRANSFER_ERROR_QUERY_URI,
            )
            .inOrder()
        assertThat(event.successFileMimeTypesCount).isEqualTo(2)
    }

    @Test
    fun transferFiles_copySuccess_doesNotRemoveOriginal() = runTest {
        val keepOriginal = true
        val testData = "file content".toByteArray()

        whenever(mockInputStream.read(any())).thenAnswer {
            val buffer = it.arguments[0] as ByteArray
            System.arraycopy(testData, 0, buffer, 0, testData.size)
            // In the next call, return -1 to signify the end of the stream
            whenever(mockInputStream.read(any())).thenReturn(-1)
            return@thenAnswer testData.size
        }

        fileTransferManagerImpl.transferFiles(uris, keepOriginal, PATH)
        val inOrder = inOrder(mockContentResolver, mockOutputStream)
        inOrder.verify(mockContentResolver).openInputStream(mockUri)
        inOrder.verify(mockContentResolver).openOutputStream(mockInsertUri)
        inOrder.verify(mockOutputStream).write(any(), eq(0), eq(testData.size))

        ExtendedMockito.verify({ DocumentsContract.deleteDocument(any(), eq(mockUri)) }, never())
    }

    @Test
    fun transferFiles_moveSuccess_removesOriginal() = runTest {
        val keepOriginal = false
        ExtendedMockito.doReturn(true).`when` { DocumentsContract.deleteDocument(any(), any()) }

        fileTransferManagerImpl.transferFiles(uris, keepOriginal, PATH)

        verify(mockContentResolver).openInputStream(mockUri)
        verify(mockContentResolver).openOutputStream(mockInsertUri)

        ExtendedMockito.verify { DocumentsContract.deleteDocument(mockContentResolver, mockUri) }
    }

    @Test
    fun transferFiles_copyFails_doesNotRemoveOriginal() = runTest {
        val keepOriginal = false
        whenever(mockInputStream.read(any())).thenThrow(IOException("Disk full"))

        fileTransferManagerImpl.transferFiles(uris, keepOriginal, PATH)

        ExtendedMockito.verify({ DocumentsContract.deleteDocument(any(), any()) }, never())
    }

    @Test
    fun transferFiles_updatesProgressNotification() = runTest {
        val totalBytes = 1000L
        val chunkSize = 100
        whenever(mockStorageHelper.calculateTotalSize(any(), any(), any())).thenReturn(totalBytes)

        // Simulate reading two chunks of data before the stream ends
        whenever(mockInputStream.read(any()))
            .thenReturn(chunkSize) // First read
            .thenReturn(chunkSize) // Second read
            .thenReturn(-1) // End of stream

        whenever(mockClock.elapsedRealtime()).thenReturn(0L).thenReturn(100L).thenReturn(1200L)

        fileTransferManagerImpl.transferFiles(uris, true, PATH)

        verify(mockNotificationsHelper).updateProgressNotification(eq(20), eq(1), eq(true))
    }

    private fun getPrivateSpaceMoveContentStatsEvent(): PrivateSpaceMoveContentEvent {
        ExtendedMockito.verify { StatsLog.write(statsEventCaptor.capture()) }
        val atom = StatsEventTestUtils.convertToAtom(statsEventCaptor.firstValue, registry)
        assertThat(atom.hasExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent))
            .isTrue()
        val event = atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)
        return event
    }
}