/*
 * Copyright (C) 2022 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.ondevicepersonalization.services.inference;

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

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;

import android.adservices.ondevicepersonalization.Constants;
import android.adservices.ondevicepersonalization.InferenceInput;
import android.adservices.ondevicepersonalization.InferenceInputParcel;
import android.adservices.ondevicepersonalization.InferenceOutput;
import android.adservices.ondevicepersonalization.InferenceOutputParcel;
import android.adservices.ondevicepersonalization.ModelId;
import android.adservices.ondevicepersonalization.OnDevicePersonalizationException;
import android.adservices.ondevicepersonalization.RemoteDataImpl;
import android.adservices.ondevicepersonalization.aidl.IDataAccessService;
import android.adservices.ondevicepersonalization.aidl.IDataAccessServiceCallback;
import android.adservices.ondevicepersonalization.aidl.IIsolatedModelServiceCallback;
import android.content.Context;
import android.net.Uri;
import android.os.Bundle;
import android.os.ParcelFileDescriptor;

import androidx.test.core.app.ApplicationProvider;

import com.android.ondevicepersonalization.internal.util.ByteArrayUtil;
import com.android.ondevicepersonalization.internal.util.LoggerFactory;

import org.junit.Before;
import org.junit.Test;

import java.io.File;
import java.io.InputStream;
import java.nio.file.Files;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class IsolatedModelServiceImplTest {
    private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
    private static final String TAG = IsolatedModelServiceImplTest.class.getSimpleName();

    private static final String MODEL_KEY = "modelKey";
    private final Random mRandom = new Random();

    private Bundle mCallbackResult;
    private InferenceInput.Params mParams;
    private RemoteDataImpl mRemoteData;

    @Before
    public void setup() {
        mRemoteData =
                new RemoteDataImpl(
                        IDataAccessService.Stub.asInterface(new TestDataAccessService()));
        mParams = new InferenceInput.Params.Builder(mRemoteData, MODEL_KEY).build();
    }

    @Test
    public void runModelInference_singleExample_success() throws Exception {
        // Set up inference context.
        InferenceInput inferenceInput =
                new InferenceInput.Builder(
                                mParams, generateInferenceInput(1), generateInferenceOutput(1))
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        InferenceOutputParcel result = verifyAndGetCallbackResult(callback);
        Map<Integer, Object> outputs =
                (Map<Integer, Object>) ByteArrayUtil.deserializeObject(result.getData());
        float[] output1 = (float[]) outputs.get(0);
        assertThat(output1.length).isEqualTo(1);
    }

    @Test
    public void runModelInference_setBatchSizeNotMatch_success() throws Exception {
        // Set up inference context.
        int numExample = 50;
        InferenceInput inferenceInput =
                new InferenceInput.Builder(
                                mParams,
                                generateInferenceInput(numExample),
                                generateInferenceOutput(numExample))
                        .setBatchSize(numExample - 10)
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        InferenceOutputParcel result = verifyAndGetCallbackResult(callback);
        Map<Integer, Object> outputs =
                (Map<Integer, Object>) ByteArrayUtil.deserializeObject(result.getData());
        float[] output1 = (float[]) outputs.get(0);
        assertThat(output1.length).isEqualTo(numExample);
    }

    @Test
    public void runModelInferenceBatch_setBatchSize_success() throws Exception {
        int numExample = 50;
        InferenceInput inferenceInput =
                new InferenceInput.Builder(
                                mParams,
                                generateInferenceInput(numExample),
                                generateInferenceOutput(numExample))
                        .setBatchSize(numExample)
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        InferenceOutputParcel result = verifyAndGetCallbackResult(callback);
        Map<Integer, Object> outputs =
                (Map<Integer, Object>) ByteArrayUtil.deserializeObject(result.getData());
        float[] output1 = (float[]) outputs.get(0);
        assertThat(output1.length).isEqualTo(numExample);
    }

    @Test
    public void runModelInferenceBatch_notSetBatchSize_success() throws Exception {
        int numExample = 50;
        InferenceInput inferenceInput =
                new InferenceInput.Builder(
                                mParams,
                                generateInferenceInput(numExample),
                                generateInferenceOutput(numExample))
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        InferenceOutputParcel result = verifyAndGetCallbackResult(callback);
        Map<Integer, Object> outputs =
                (Map<Integer, Object>) ByteArrayUtil.deserializeObject(result.getData());
        float[] output1 = (float[]) outputs.get(0);
        assertThat(output1.length).isEqualTo(numExample);
    }

    @Test
    public void runModelInference_invalidInputFormat() throws Exception {
        // Misconfigured inputs.
        float input0 = 1.23f;
        float input1 = 2.43f;
        Object[] invalidInput = {input0, input1, input0};

        InferenceInput inferenceInput =
                new InferenceInput.Builder(mParams, invalidInput, generateInferenceOutput(1))
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        verifyCallBackError(callback, OnDevicePersonalizationException.ERROR_INFERENCE_FAILED);
    }

    @Test
    public void missModelId_throws() {
        assertThrows(
                NullPointerException.class,
                () ->
                        new InferenceInput.Builder(
                                        new InferenceInput.Params.Builder(mRemoteData, null)
                                                .build(),
                                        generateInferenceInput(1),
                                        generateInferenceOutput(1))
                                .build());
    }

    @Test
    public void missingInputData_buildThrows() {
        assertThrows(
                NullPointerException.class,
                () ->
                        new InferenceInput.Builder(mParams, null, generateInferenceOutput(1))
                                .build());
    }

    @Test
    public void negativeThreadNum_buildThrows() {
        assertThrows(
                IllegalStateException.class,
                () ->
                        new InferenceInput.Params.Builder(mRemoteData, MODEL_KEY)
                                .setRecommendedNumThreads(-1)
                                .build());
    }

    @Test
    public void runModelInference_modelNotExist() throws Exception {
        InferenceInput.Params params =
                new InferenceInput.Params.Builder(mRemoteData, "nonexist").build();
        InferenceInput inferenceInput =
                // Not set output structure in InferenceOutput.
                new InferenceInput.Builder(
                                params, generateInferenceInput(1), generateInferenceOutput(1))
                        .build();

        Bundle bundle = new Bundle();
        bundle.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        var callback = new TestServiceCallback();
        modelService.runInference(bundle, callback);

        verifyCallBackError(
                callback, OnDevicePersonalizationException.ERROR_INFERENCE_MODEL_NOT_FOUND);
    }

    @Test
    public void runModelInference_missingDataAccessServiceBinder() {
        // Set up inference context.
        InferenceInput inferenceInput =
                new InferenceInput.Builder(
                                mParams, generateInferenceInput(1), generateInferenceOutput(1))
                        .build();

        Bundle bundle = new Bundle();
        bundle.putParcelable(
                Constants.EXTRA_INFERENCE_INPUT, new InferenceInputParcel(inferenceInput));

        IsolatedModelServiceImpl modelService = new IsolatedModelServiceImpl();
        assertThrows(
                NullPointerException.class,
                () -> modelService.runInference(bundle, new TestServiceCallback()));
    }

    private void verifyCallBackError(TestServiceCallback callback, int errorCode) throws Exception {
        assertTrue(
                "Timeout when run ModelService.runInference complete",
                callback.mLatch.await(5000, TimeUnit.SECONDS));
        assertTrue(callback.mError);
        assertThat(callback.mErrorCode).isEqualTo(errorCode);
    }

    private InferenceOutputParcel verifyAndGetCallbackResult(TestServiceCallback callback)
            throws Exception {
        assertTrue(
                "Timeout when run ModelService.runInference complete",
                callback.mLatch.await(5000, TimeUnit.SECONDS));
        assertFalse(callback.mError);
        return mCallbackResult.getParcelable(Constants.EXTRA_RESULT, InferenceOutputParcel.class);
    }

    private Object[] generateInferenceInput(int numExample) {
        int numFloatFeatures = 13;
        int numStringFeatures = 26;
        int totalInputs = numFloatFeatures + numStringFeatures;

        Object[] inputs = new Object[totalInputs];
        // Generate 13 float input tensors
        for (int i = 0; i < numFloatFeatures; i++) {
            float[][] floatInput = new float[numExample][1];
            for (int j = 0; j < numExample; j++) {
                floatInput[j][0] = mRandom.nextFloat() * 100;
            }
            inputs[i] = floatInput;
        }
        // Generate 26 string input tensors
        for (int i = 0; i < numStringFeatures; i++) {
            String[][] stringInput = new String[numExample][1];
            for (int j = 0; j < numExample; j++) {
                stringInput[j][0] = "example_cat_" + mRandom.nextInt(100);
            }
            inputs[numFloatFeatures + i] = stringInput;
        }
        return inputs;
    }

    private InferenceOutput generateInferenceOutput(int numExample) {
        float[] output0 = new float[numExample];
        HashMap<Integer, Object> outputMap = new HashMap<>();
        outputMap.put(0, output0);
        return new InferenceOutput.Builder().setDataOutputs(outputMap).build();
    }

    class TestServiceCallback extends IIsolatedModelServiceCallback.Stub {
        public boolean mError = false;
        public int mErrorCode = 0;
        private final CountDownLatch mLatch = new CountDownLatch(1);

        @Override
        public void onSuccess(Bundle result) {
            mCallbackResult = result;
            mLatch.countDown();
        }

        @Override
        public void onError(int errorCode) {
            mError = true;
            mErrorCode = errorCode;
            mLatch.countDown();
        }
    }

    static class TestDataAccessService extends IDataAccessService.Stub {
        @Override
        public void onRequest(int operation, Bundle params, IDataAccessServiceCallback callback) {
            if (operation == Constants.DATA_ACCESS_OP_GET_MODEL) {
                try {
                    ModelId modelId = params.getParcelable(Constants.EXTRA_MODEL_ID, ModelId.class);
                    sLogger.d(
                            TAG
                                    + " TestDataAccessService onRequest model id %s"
                                    + modelId.getKey());
                    if (modelId.getKey().equals("nonexist")) {
                        callback.onSuccess(Bundle.EMPTY);
                        return;
                    }
                    assertThat(modelId.getKey()).isEqualTo(MODEL_KEY);
                    assertThat(modelId.getTableId()).isEqualTo(ModelId.TABLE_ID_REMOTE_DATA);
                    Context context = ApplicationProvider.getApplicationContext();
                    Uri modelUri =
                            Uri.parse(
                                    "android.resource://com.android.ondevicepersonalization.servicetests/raw/model");
                    File modelFile = File.createTempFile("model", ".tflite");
                    InputStream in = context.getContentResolver().openInputStream(modelUri);
                    Files.copy(in, modelFile.toPath(), REPLACE_EXISTING);
                    in.close();
                    ParcelFileDescriptor inputModelFd =
                            ParcelFileDescriptor.open(
                                    modelFile, ParcelFileDescriptor.MODE_READ_ONLY);

                    Bundle bundle = new Bundle();
                    bundle.putParcelable(Constants.EXTRA_RESULT, inputModelFd);
                    callback.onSuccess(bundle);
                } catch (Exception e) {
                    // Ignore
                }
            }
        }

        @Override
        public void logApiCallStats(int apiName, long latencyMillis, int responseCode) {}
    }
}