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

import android.multiuser.Flags
import android.os.statsd.privatespace.FileSizeBucket
import android.os.statsd.privatespace.MimeTypeCategory
import android.os.statsd.privatespace.NotificationEventType
import android.os.statsd.privatespace.OperationType
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.util.StatsEvent
import android.util.StatsEventTestUtils
import android.util.StatsLog
import androidx.test.ext.junit.runners.AndroidJUnit4
import com.android.dx.mockito.inline.extended.ExtendedMockito
import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
import com.android.privatespace.PrivatespaceExtensionAtoms
import com.google.common.truth.Truth.assertThat
import com.google.protobuf.ExtensionRegistryLite
import org.junit.After
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

@RequiresFlagsEnabled(Flags.FLAG_ENABLE_MOVING_CONTENT_INTO_PRIVATE_SPACE)
@RunWith(AndroidJUnit4::class)
class FileTransferResultLoggerTest {

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

    @Captor val statsEventCaptor = argumentCaptor<StatsEvent>()

    private val registry = ExtensionRegistryLite.newInstance()

    private lateinit var mockitoSession: MockitoSession

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

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

    @Test
    fun logMoveContentResultEvent_success_singleFile() {
        val filesLoggingDetails = FilesLoggingDetails()
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 1111, mimeType = "image/jpeg")

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = false,
            filesSelectedCount = 1,
            filesLoggingDetails = filesLoggingDetails,
        )

        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)

        assertThat(event.operationType).isEqualTo(OperationType.OPERATION_MOVE)
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_SUCCESS)
        assertThat(event.filesSelectedCount).isEqualTo(1)
        assertThat(event.filesSuccessCount).isEqualTo(1)
        assertThat(event.totalSuccessFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_1_TO_10_KB)
        assertThat(event.successFileMimeTypesList)
            .containsExactly(MimeTypeCategory.MIME_TYPE_CATEGORY_IMAGE)

        // Assert other fields are not set
        assertThat(event.filesFailureCount).isEqualTo(0)
        assertThat(event.totalFailureFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.failedFileErrorCodesCount).isEqualTo(0)
        assertThat(event.failedFileMimeTypesCount).isEqualTo(0)
        assertThat(event.notificationEvent)
            .isEqualTo(NotificationEventType.NOTIFICATION_EVENT_TYPE_UNKNOWN)
        assertThat(event.durationMs).isEqualTo(0)
    }

    @Test
    fun logMoveContentResultEvent_success_multipleFiles() {
        val filesLoggingDetails = FilesLoggingDetails()
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 2000, mimeType = "image/png")
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 15000, mimeType = "video/mp4")
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 300000, mimeType = "audio/mp3")

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = true,
            filesSelectedCount = 3,
            filesLoggingDetails = filesLoggingDetails,
        )

        ExtendedMockito.verify { StatsLog.write(statsEventCaptor.capture()) }
        val atom = StatsEventTestUtils.convertToAtom(statsEventCaptor.firstValue, registry)
        val event = atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)

        assertThat(event.operationType).isEqualTo(OperationType.OPERATION_COPY)
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_SUCCESS)
        assertThat(event.filesSelectedCount).isEqualTo(3)
        assertThat(event.filesSuccessCount).isEqualTo(3)
        assertThat(event.totalSuccessFilesSizeBucket)
            .isEqualTo(
                FileSizeBucket.FILE_SIZE_100_TO_500_KB
            ) // 2000 + 15000 + 300000 = 317000 bytes
        assertThat(event.successFileMimeTypesList)
            .containsExactly(
                MimeTypeCategory.MIME_TYPE_CATEGORY_IMAGE,
                MimeTypeCategory.MIME_TYPE_CATEGORY_VIDEO,
                MimeTypeCategory.MIME_TYPE_CATEGORY_AUDIO,
            )
            .inOrder()

        // Assert other fields are not set
        assertThat(event.filesFailureCount).isEqualTo(0)
        assertThat(event.totalFailureFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.failedFileErrorCodesCount).isEqualTo(0)
        assertThat(event.failedFileMimeTypesCount).isEqualTo(0)
        assertThat(event.notificationEvent)
            .isEqualTo(NotificationEventType.NOTIFICATION_EVENT_TYPE_UNKNOWN)
        assertThat(event.durationMs).isEqualTo(0)
    }

    @Test
    fun logMoveContentResultEvent_partialSuccess() {
        val filesLoggingDetails = FilesLoggingDetails()
        // Success
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 5000, mimeType = "image/gif")
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 25000, mimeType = "application/pdf")
        // Failure
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE_VALUE,
            sizeBytes = 100000,
            mimeType = "video/quicktime",
        )

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = false,
            filesSelectedCount = 3, // 2 success + 1 failure
            filesLoggingDetails = filesLoggingDetails,
        )

        ExtendedMockito.verify { StatsLog.write(statsEventCaptor.capture()) }
        val atom = StatsEventTestUtils.convertToAtom(statsEventCaptor.firstValue, registry)
        val event = atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)

        // Assert primary fields
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_PARTIAL_SUCCESS)
        assertThat(event.filesSelectedCount).isEqualTo(3)

        // Assert success fields
        assertThat(event.filesSuccessCount).isEqualTo(2)
        assertThat(event.totalSuccessFilesSizeBucket)
            .isEqualTo(FileSizeBucket.FILE_SIZE_10_TO_100_KB) // 5000 + 25000 = 30000 bytes
        assertThat(event.successFileMimeTypesList)
            .containsExactly(
                MimeTypeCategory.MIME_TYPE_CATEGORY_IMAGE,
                MimeTypeCategory.MIME_TYPE_CATEGORY_PDF,
            )
            .inOrder()

        // Assert failure fields
        assertThat(event.filesFailureCount).isEqualTo(1)
        assertThat(event.totalFailureFilesSizeBucket)
            .isEqualTo(FileSizeBucket.FILE_SIZE_10_TO_100_KB)
        assertThat(event.failedFileErrorCodesCount).isEqualTo(1)
        assertThat(event.failedFileErrorCodesList[0])
            .isEqualTo(TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE)
        assertThat(event.failedFileMimeTypesCount).isEqualTo(1)
        assertThat(event.failedFileMimeTypesList[0])
            .isEqualTo(MimeTypeCategory.MIME_TYPE_CATEGORY_VIDEO)
    }

    @Test
    fun logMoveContentResultEvent_partialSuccess_checkOrderOfMimeTypesAndErrorCodes() {
        val filesLoggingDetails = FilesLoggingDetails()
        // Success
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 5000, mimeType = "image/gif")
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 25000, mimeType = "application/pdf")
        // Failure
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_COPY_FILE_VALUE,
            sizeBytes = 100000,
            mimeType = "video/quicktime",
        )
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE_VALUE,
            sizeBytes = 100000,
            mimeType = "audio/mp3",
        )
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE_VALUE,
            sizeBytes = 100000,
            mimeType = "video/mp4",
        )

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = true,
            filesSelectedCount = 5, // 2 success + 3 failure
            filesLoggingDetails = filesLoggingDetails,
        )

        ExtendedMockito.verify { StatsLog.write(statsEventCaptor.capture()) }
        val atom = StatsEventTestUtils.convertToAtom(statsEventCaptor.firstValue, registry)
        val event = atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)

        // Assert primary fields
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_PARTIAL_SUCCESS)
        assertThat(event.filesSelectedCount).isEqualTo(5)

        // Assert success fields
        assertThat(event.filesSuccessCount).isEqualTo(2)
        assertThat(event.successFileMimeTypesList)
            .containsExactly(
                MimeTypeCategory.MIME_TYPE_CATEGORY_IMAGE,
                MimeTypeCategory.MIME_TYPE_CATEGORY_PDF,
            )
            .inOrder()

        // Assert failure fields
        assertThat(event.filesFailureCount).isEqualTo(3)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(
                TransferErrorCode.TRANSFER_ERROR_COPY_FILE,
                TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE,
                TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE,
            )
            .inOrder()
        assertThat(event.failedFileMimeTypesList)
            .containsExactly(
                MimeTypeCategory.MIME_TYPE_CATEGORY_VIDEO,
                MimeTypeCategory.MIME_TYPE_CATEGORY_AUDIO,
                MimeTypeCategory.MIME_TYPE_CATEGORY_VIDEO,
            )
            .inOrder()
    }

    @Test
    fun logMoveContentResultEvent_partialSuccess_withNullMimeType() {
        val filesLoggingDetails = FilesLoggingDetails()
        filesLoggingDetails.addSuccessFileDetails(sizeBytes = 5000, mimeType = "image/gif")
        // Failed file with a null mime type
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_GET_METADATA_VALUE,
            sizeBytes = 200000,
            mimeType = null,
        )

        // Failed file without mime type
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE_VALUE,
            sizeBytes = 300000,
        )

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = false,
            filesSelectedCount = 2,
            filesLoggingDetails = filesLoggingDetails,
        )

        ExtendedMockito.verify { StatsLog.write(statsEventCaptor.capture()) }
        val atom = StatsEventTestUtils.convertToAtom(statsEventCaptor.firstValue, registry)
        val event = atom.getExtension(PrivatespaceExtensionAtoms.privateSpaceMoveContentEvent)

        // Assert primary fields
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_PARTIAL_SUCCESS)
        assertThat(event.filesSelectedCount).isEqualTo(2)

        // Assert success fields
        assertThat(event.filesSuccessCount).isEqualTo(1)
        assertThat(event.successFileMimeTypesList).isNotEmpty()

        // Assert failure fields
        assertThat(event.filesFailureCount).isEqualTo(2)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(
                TransferErrorCode.TRANSFER_ERROR_GET_METADATA,
                TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE,
            )
            .inOrder()
        assertThat(event.failedFileMimeTypesList)
            .containsExactly(
                MimeTypeCategory.MIME_TYPE_CATEGORY_UNKNOWN,
                MimeTypeCategory.MIME_TYPE_CATEGORY_UNKNOWN,
            )
    }

    @Test
    fun logMoveContentResultEvent_onlyFailedFiles() {
        val filesLoggingDetails = FilesLoggingDetails()
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_CALCULATE_FILE_SIZE_VALUE,
            sizeBytes = 100_000,
            mimeType = "image/png",
        )
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_NOT_ENOUGH_SPACE_VALUE,
            sizeBytes = 100_000,
            mimeType = "application/pdf",
        )
        filesLoggingDetails.addFailedFileDetails(
            errorCode = TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE_VALUE,
            sizeBytes = 99_999,
            mimeType = "video/mp4",
        )
        val fileSelectedCount = 3

        FileTransferResultLogger.logMoveContentResultEvent(
            keepOriginal = false,
            filesSelectedCount = fileSelectedCount,
            filesLoggingDetails = filesLoggingDetails,
        )

        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)

        assertThat(event.operationType).isEqualTo(OperationType.OPERATION_MOVE)
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.filesSelectedCount).isEqualTo(fileSelectedCount)
        assertThat(event.filesFailureCount).isEqualTo(fileSelectedCount)
        assertThat(event.totalFailureFilesSizeBucket)
            .isEqualTo(FileSizeBucket.FILE_SIZE_100_TO_500_KB)

        // Assert other fields are not set
        assertThat(event.filesSuccessCount).isEqualTo(0)
        assertThat(event.totalSuccessFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.failedFileErrorCodesCount).isEqualTo(3)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(
                TransferErrorCode.TRANSFER_ERROR_CALCULATE_FILE_SIZE,
                TransferErrorCode.TRANSFER_ERROR_NOT_ENOUGH_SPACE,
                TransferErrorCode.TRANSFER_ERROR_REMOVE_ORIGINAL_FILE,
            )
        assertThat(event.failedFileMimeTypesCount).isEqualTo(3)
        assertThat(event.notificationEvent)
            .isEqualTo(NotificationEventType.NOTIFICATION_EVENT_TYPE_UNKNOWN)
        assertThat(event.durationMs).isEqualTo(0)
    }

    @Test
    fun logTransferFailedOnPrechecks() {
        val errorCode = TransferErrorCode.TRANSFER_ERROR_ABOVE_AVAILABLE_DEVICE_STORAGE
        val numberOfFiles = 3
        val totalSizeFilesSelected = 50000L // 50KB

        FileTransferResultLogger.logTransferFailedOnPrechecks(
            errorCode.number,
            numberOfFiles,
            totalSizeFilesSelected,
        )

        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)

        assertThat(event.operationType).isEqualTo(OperationType.OPERATION_UNKNOWN)
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.filesSelectedCount).isEqualTo(numberOfFiles)
        assertThat(event.filesFailureCount).isEqualTo(numberOfFiles)
        assertThat(event.totalFailureFilesSizeBucket)
            .isEqualTo(FileSizeBucket.FILE_SIZE_10_TO_100_KB)
        assertThat(event.failedFileErrorCodesList).isEqualTo(List(numberOfFiles) { errorCode })

        // Assert other fields are not set
        assertThat(event.filesSuccessCount).isEqualTo(0)
        assertThat(event.totalSuccessFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.failedFileMimeTypesCount).isEqualTo(0)
        assertThat(event.successFileMimeTypesList).isEmpty()
        assertThat(event.notificationEvent)
            .isEqualTo(NotificationEventType.NOTIFICATION_EVENT_TYPE_UNKNOWN)
        assertThat(event.durationMs).isEqualTo(0)
    }

    @Test
    fun logAnotherTransferInProgressEvent() {
        FileTransferResultLogger.logAnotherTransferInProgressEvent()

        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)

        assertThat(event.operationType).isEqualTo(OperationType.OPERATION_UNKNOWN)
        assertThat(event.transferResult).isEqualTo(TransferResult.TRANSFER_RESULT_FAILURE)
        assertThat(event.failedFileErrorCodesList)
            .containsExactly(TransferErrorCode.TRANSFER_ERROR_ANOTHER_TRANSFER_IN_PROGRESS)

        // Assert other fields are not set or are 0
        assertThat(event.filesSelectedCount).isEqualTo(0)
        assertThat(event.filesSuccessCount).isEqualTo(0)
        assertThat(event.filesFailureCount).isEqualTo(0)
        assertThat(event.totalSuccessFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.totalFailureFilesSizeBucket).isEqualTo(FileSizeBucket.FILE_SIZE_UNKNOWN)
        assertThat(event.failedFileMimeTypesCount).isEqualTo(0)
        assertThat(event.successFileMimeTypesList).isEmpty()
        assertThat(event.notificationEvent)
            .isEqualTo(NotificationEventType.NOTIFICATION_EVENT_TYPE_UNKNOWN)
        assertThat(event.durationMs).isEqualTo(0)
    }
}