/*
 * Copyright (C) 2023 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.federatedcompute.services.encryption;

import static com.android.adservices.shared.spe.JobServiceConstants.SCHEDULING_RESULT_CODE_FAILED;
import static com.android.adservices.shared.spe.JobServiceConstants.SCHEDULING_RESULT_CODE_SKIPPED;
import static com.android.adservices.shared.spe.JobServiceConstants.SCHEDULING_RESULT_CODE_SUCCESSFUL;
import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
import static com.android.odp.module.common.encryption.OdpEncryptionKey.KEY_TYPE_ENCRYPTION;

import static com.google.common.truth.Truth.assertThat;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

import android.app.job.JobInfo;
import android.app.job.JobParameters;
import android.app.job.JobScheduler;
import android.content.Context;

import androidx.test.core.app.ApplicationProvider;

import com.android.dx.mockito.inline.extended.ExtendedMockito;
import com.android.federatedcompute.services.common.FederatedComputeExecutors;
import com.android.federatedcompute.services.common.FederatedComputeJobInfo;
import com.android.federatedcompute.services.common.FlagsFactory;
import com.android.federatedcompute.services.common.PhFlagsTestUtil;
import com.android.federatedcompute.services.data.FederatedComputeDbHelper;
import com.android.federatedcompute.services.data.FederatedComputeEncryptionKeyDaoUtils;
import com.android.federatedcompute.services.sharedlibrary.spe.FederatedComputeJobScheduler;
import com.android.odp.module.common.EventLogger;
import com.android.odp.module.common.MonotonicClock;
import com.android.odp.module.common.data.OdpEncryptionKeyDao;
import com.android.odp.module.common.encryption.OdpEncryptionKey;
import com.android.odp.module.common.encryption.OdpEncryptionKeyManager;
import com.android.odp.module.common.http.HttpClient;

import com.google.common.util.concurrent.FluentFuture;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.mockito.MockitoSession;
import org.mockito.quality.Strictness;

import java.util.List;
import java.util.Optional;
import java.util.concurrent.ExecutionException;

// TODO: add tests with Ph flags
@RunWith(JUnit4.class)
public class BackgroundKeyFetchJobServiceTest {

    private BackgroundKeyFetchJobService mSpyService;

    private MockitoSession mStaticMockSession;

    private Context mContext;

    private JobScheduler mJobScheduler;

    private HttpClient mHttpClient;

    public OdpEncryptionKeyDao mEncryptionDao;

    public OdpEncryptionKeyManager mSpyKeyManager;

    private TestInjector mInjector;

    private FederatedComputeDbHelper mTestDbHelper;

    @Mock
    private EventLogger mMockEventLogger;

    @Before
    public void setUp() throws Exception {
        PhFlagsTestUtil.setUpDeviceConfigPermissions();
        PhFlagsTestUtil.disableGlobalKillSwitch();
        PhFlagsTestUtil.enableScheduleBackgroundKeyFetchJob();
        PhFlagsTestUtil.disableSpeBackgroundKeyFetchJob();
        MockitoAnnotations.initMocks(this);

        mContext = ApplicationProvider.getApplicationContext();
        mInjector = new TestInjector();
        mTestDbHelper = FederatedComputeDbHelper.getNonSingletonInstanceForTest(mContext);
        mEncryptionDao =
                FederatedComputeEncryptionKeyDaoUtils.getInstanceForTest(mContext, mTestDbHelper);
        mHttpClient = new HttpClient(/* retryLimit= */ 3, MoreExecutors.newDirectExecutorService());
        mSpyService = spy(new BackgroundKeyFetchJobService(new TestInjector()));
        doReturn(mSpyService).when(mSpyService).getApplicationContext();
        doNothing().when(mSpyService).jobFinished(any(), anyBoolean());
        mJobScheduler = mContext.getSystemService(JobScheduler.class);
        mJobScheduler.cancel(FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);
        mSpyKeyManager =
                spy(
                        FederatedComputeEncryptionKeyManagerUtils.getInstanceForTest(
                                MonotonicClock.getInstance(),
                                mEncryptionDao,
                                FlagsFactory.getFlags(),
                                mHttpClient,
                                MoreExecutors.newDirectExecutorService(),
                                mTestDbHelper));
        mStaticMockSession =
                ExtendedMockito.mockitoSession()
                        .mockStatic(FederatedComputeJobScheduler.class)
                        .initMocks(this)
                        .strictness(Strictness.LENIENT)
                        .startMocking();
    }

    @After
    public void tearDown() {
        if (mStaticMockSession != null) {
            mStaticMockSession.finishMocking();
        }

        mTestDbHelper.getWritableDatabase().close();
        mTestDbHelper.getReadableDatabase().close();
        mTestDbHelper.close();
    }

    @Test
    public void testDefaultNoArgConstructor() {
        BackgroundKeyFetchJobService instance = new BackgroundKeyFetchJobService();
        assertNotNull("default no-arg constructor is required by JobService", instance);
    }

    @Test
    public void testOnStartJob() {
        OdpEncryptionKeyManager keyManager = mInjector.getEncryptionKeyManager(mContext);
        List<OdpEncryptionKey> emptyKeyList = List.of();
        doReturn(FluentFuture.from(Futures.immediateFuture(emptyKeyList)))
                .when(keyManager)
                .fetchAndPersistActiveKeys(
                        KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true, Optional.empty());

        mSpyService.run(mock(JobParameters.class));

        verify(mSpyService, times(1)).onStartJob(any());
        verify(mSpyService, times(1)).jobFinished(any(), anyBoolean());
        verify(mMockEventLogger, times(1)).logEncryptionKeyFetchStartEventKind();
    }

    @Test
    public void onStartJobTestSpeEnabled() {
        PhFlagsTestUtil.enableSpeBackgroundKeyFetchJob();

        // Mock OdpJobScheduler to not actually schedule the job.
        FederatedComputeJobScheduler mockedScheduler = mock(FederatedComputeJobScheduler.class);
        doReturn(mockedScheduler).when(() -> FederatedComputeJobScheduler.getInstance(any()));

        assertThat(mSpyService.onStartJob(mock(JobParameters.class))).isFalse();

        // Verify SPE scheduler has rescheduled the job.
        verify(mockedScheduler).schedule(any(), any());
    }

    @Test
    public void testOnStartJob_onFailure() {
        OdpEncryptionKeyManager keyManager = mInjector.getEncryptionKeyManager(mContext);
        doReturn(
                        FluentFuture.from(
                                Futures.immediateFailedFuture(
                                        new ExecutionException(
                                                " Failed to fetch keys",
                                                new IllegalStateException("http 404")))))
                .when(keyManager)
                .fetchAndPersistActiveKeys(
                        KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true, Optional.empty());

        mSpyService.run(mock(JobParameters.class));

        verify(mSpyService, times(1)).onStartJob(any());
        verify(mSpyService, times(1)).jobFinished(any(), anyBoolean());
        verify(mMockEventLogger, times(1)).logEncryptionKeyFetchStartEventKind();
    }

    @Test
    public void testScheduleJob() {
        final JobScheduler jobScheduler = mContext.getSystemService(JobScheduler.class);

        assertThat(
                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                                mContext, FlagsFactory.getFlags(), /* forceSchedule */ false))
                .isEqualTo(SCHEDULING_RESULT_CODE_SUCCESSFUL);

        final JobInfo scheduledJob =
                jobScheduler.getPendingJob(
                        FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);

        assertThat(scheduledJob.getId())
                .isEqualTo(FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);
    }

    @Test
    public void testScheduleJob_notNeeded() {
        assertThat(
                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                                mContext, FlagsFactory.getFlags(), /* forceSchedule */ false))
                .isEqualTo(SCHEDULING_RESULT_CODE_SUCCESSFUL);

        assertThat(
                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                                mContext, FlagsFactory.getFlags(), /* forceSchedule */ false))
                .isEqualTo(SCHEDULING_RESULT_CODE_SKIPPED);
    }

    @Test
    public void testScheduleJob_forceScheduleSuccessful() {
        assertThat(
                BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                        mContext, FlagsFactory.getFlags(), /* forceSchedule */ false))
                .isEqualTo(SCHEDULING_RESULT_CODE_SUCCESSFUL);

        assertThat(
                BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                        mContext, FlagsFactory.getFlags(), /* forceSchedule */ true))
                .isEqualTo(SCHEDULING_RESULT_CODE_SUCCESSFUL);
    }

    @Test
    public void testScheduleJob_notEnabled() {
        PhFlagsTestUtil.disableScheduleBackgroundKeyFetchJob();

        assertThat(BackgroundKeyFetchJobService.scheduleJobIfNeeded(
                mContext, FlagsFactory.getFlags(), /* forceSchedule */ false
        )).isEqualTo(SCHEDULING_RESULT_CODE_FAILED);
    }

    @Test
    public void testOnStopJob() {
        assertFalse(mSpyService.onStopJob(mock(JobParameters.class)));
    }

    @Test
    public void testOnStartJob_enableKillSwitch() {
        PhFlagsTestUtil.enableGlobalKillSwitch();
        OdpEncryptionKeyManager keyManager = mInjector.getEncryptionKeyManager(mContext);
        List<OdpEncryptionKey> emptyKeyList = List.of();
        doReturn(FluentFuture.from(Futures.immediateFuture(emptyKeyList)))
                .when(keyManager)
                .fetchAndPersistActiveKeys(
                        KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true, Optional.empty());
        doReturn(mJobScheduler).when(mSpyService).getSystemService(JobScheduler.class);
        mSpyService
                .scheduleJobIfNeeded(mContext, FlagsFactory.getFlags(), /* forceSchedule */ false);
        assertTrue(mJobScheduler.getPendingJob(
                FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID)
                != null);
        doNothing().when(mSpyService).jobFinished(any(), anyBoolean());
        boolean result = mSpyService.onStartJob(mock(JobParameters.class));
        assertTrue(result);
        verify(mSpyService, times(1)).jobFinished(any(), eq(false));
        verify(keyManager, never()).fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION,
                /* isScheduledJob= */ true, Optional.empty());
        assertTrue(mJobScheduler.getPendingJob(
                FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID)
                == null);
    }

    @Test
    public void testDefaultInjector() {
        BackgroundKeyFetchJobService.Injector injector =
                new BackgroundKeyFetchJobService.Injector();
        assertThat(injector.getExecutor())
                .isEqualTo(FederatedComputeExecutors.getBackgroundExecutor());
        assertThat(injector.getEncryptionKeyManager(mContext))
                .isEqualTo(FederatedComputeEncryptionKeyManagerUtils.getInstance(mContext));
        assertThat(injector.getLightWeightExecutor())
                .isEqualTo(FederatedComputeExecutors.getLightweightExecutor());
    }

    class TestInjector extends BackgroundKeyFetchJobService.Injector {
        @Override
        ListeningExecutorService getExecutor() {
            return MoreExecutors.newDirectExecutorService();
        }

        @Override
        ListeningExecutorService getLightWeightExecutor() {
            return MoreExecutors.newDirectExecutorService();
        }

        @Override
        OdpEncryptionKeyManager getEncryptionKeyManager(Context context) {
            return mSpyKeyManager;
        }

        @Override
        EventLogger getEventLogger() {
            return mMockEventLogger;
        }
    }
}