/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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 org.apache.atlas.repository.impexp;

import org.apache.atlas.RequestContext;
import org.apache.atlas.TestModules;
import org.apache.atlas.TestUtilsV2;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.impexp.AtlasExportRequest;
import org.apache.atlas.model.impexp.AtlasExportResult;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.AtlasTestBase;
import org.apache.atlas.repository.graph.AtlasGraphProvider;
import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
import org.apache.atlas.repository.store.graph.v2.AtlasEntityStoreV2;
import org.apache.atlas.repository.store.graph.v2.AtlasEntityStream;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.Guice;
import org.testng.annotations.Test;

import javax.inject.Inject;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;

@Guice(modules = TestModules.TestOnlyModule.class)
public class ExportServiceTest extends AtlasTestBase {
    private static final Logger LOG = LoggerFactory.getLogger(ExportServiceTest.class);

    @Inject
    AtlasTypeRegistry typeRegistry;
    @Inject
    ExportService exportService;
    @Inject
    private AtlasTypeDefStore typeDefStore;
    @Inject
    private ExportImportAuditService auditService;

    @Inject
    private AtlasEntityStoreV2 entityStore;

    @BeforeTest
    public void setupTest() throws IOException, AtlasBaseException {
        RequestContext.clear();
        RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);
        basicSetup(typeDefStore, typeRegistry);
    }

    @BeforeClass
    public void setupSampleData() throws Exception {
        super.initialize();

        AtlasTypesDef sampleTypes   = TestUtilsV2.defineDeptEmployeeTypes();
        AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(sampleTypes, typeRegistry);

        if (!typesToCreate.isEmpty()) {
            typeDefStore.createTypesDef(typesToCreate);
        }

        AtlasEntity.AtlasEntitiesWithExtInfo deptEg2      = TestUtilsV2.createDeptEg2();
        AtlasEntityStream                    entityStream = new AtlasEntityStream(deptEg2);
        EntityMutationResponse               emr          = entityStore.createOrUpdate(entityStream, false);
        assertNotNull(emr);
        assertNotNull(emr.getCreatedEntities());
        assertFalse(emr.getCreatedEntities().isEmpty());
    }

    @AfterClass
    public void clear() throws Exception {
        Thread.sleep(1000);
        assertExportImportAuditEntry(auditService);
        AtlasGraphProvider.cleanup();

        super.cleanup();
    }

    @Test
    public void exportType() throws AtlasBaseException {
        String requestingIP = "1.0.0.0";
        String hostName     = "root";

        AtlasExportRequest    request = getRequestForFullFetch();
        ByteArrayOutputStream baos    = new ByteArrayOutputStream();
        ZipSink               zipSink = new ZipSink(baos);
        AtlasExportResult     result  = exportService.run(zipSink, request, "admin", hostName, requestingIP);

        assertNotNull(exportService);
        assertEquals(result.getHostName(), hostName);
        assertEquals(result.getClientIpAddress(), requestingIP);
        assertEquals(request, result.getRequest());
        assertNotNull(result.getSourceClusterName());
    }

    @Test(expectedExceptions = AtlasBaseException.class)
    public void requestingEntityNotFound_NoData() throws AtlasBaseException, IOException {
        String requestingIP = "1.0.0.0";
        String hostName     = "root";

        ByteArrayOutputStream baos    = new ByteArrayOutputStream();
        ZipSink               zipSink = new ZipSink(baos);
        AtlasExportResult result = exportService.run(
                zipSink, getRequestForFullFetch(), "admin", hostName, requestingIP);

        assertNull(result.getData());

        ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
        new ZipSource(bais);
    }

    @Test
    public void requestingEntityFoundDefaultFetch_ContainsData() throws Exception {
        ZipSource source = runExportWithParameters(
                getRequestForDept(false, "", false, ""));
        verifyExportForHrData(source);
    }

    @Test
    public void requestingHrEntityWithMatchTypeContains_ContainsData() throws Exception {
        ZipSource source = runExportWithParameters(
                getRequestForDept(false, "", true, "CONTAINS"));
        verifyExportForHrData(source);
    }

    @Test
    public void requestingHrEntityWithMatchTypeEndsWith_ContainsData() throws Exception {
        ZipSource source = runExportWithParameters(
                getRequestForDept(false, "", true, "ENDSWITH"));
        verifyExportForHrData(source);
    }

    @Test
    public void requestingDeptEntityFoundFullFetch_ContainsData() throws Exception {
        ZipSource source = runExportWithParameters(
                getRequestForDept(true, "FULL", false, ""));
        verifyExportForHrData(source);
    }

    @Test
    public void requestingDeptEntityFoundConnectedFetch_ContainsData() throws Exception {
        ZipSource source = runExportWithParameters(
                getRequestForDept(true, "CONNECTED", false, ""));
        verifyExportForHrDataForConnected(source);
    }

    @Test
    public void requestingEmployeeEntityFoundConnectedFetch_ContainsData() throws Exception {
        ZipSource zipSource = runExportWithParameters(getRequestForEmployee());
        verifyExportForEmployeeData(zipSource);
    }

    @Test
    public void verifyOverallStatus() {
        assertEquals(exportService.getOverallOperationStatus(), AtlasExportResult.OperationStatus.FAIL);

        assertEquals(exportService.getOverallOperationStatus(AtlasExportResult.OperationStatus.SUCCESS), AtlasExportResult.OperationStatus.SUCCESS);

        assertEquals(exportService.getOverallOperationStatus(
                AtlasExportResult.OperationStatus.SUCCESS,
                AtlasExportResult.OperationStatus.SUCCESS,
                AtlasExportResult.OperationStatus.SUCCESS), AtlasExportResult.OperationStatus.SUCCESS);

        assertEquals(exportService.getOverallOperationStatus(
                AtlasExportResult.OperationStatus.FAIL,
                AtlasExportResult.OperationStatus.PARTIAL_SUCCESS,
                AtlasExportResult.OperationStatus.SUCCESS), AtlasExportResult.OperationStatus.PARTIAL_SUCCESS);

        assertEquals(exportService.getOverallOperationStatus(
                AtlasExportResult.OperationStatus.FAIL,
                AtlasExportResult.OperationStatus.FAIL,
                AtlasExportResult.OperationStatus.PARTIAL_SUCCESS), AtlasExportResult.OperationStatus.PARTIAL_SUCCESS);

        assertEquals(exportService.getOverallOperationStatus(
                AtlasExportResult.OperationStatus.FAIL,
                AtlasExportResult.OperationStatus.FAIL,
                AtlasExportResult.OperationStatus.FAIL), AtlasExportResult.OperationStatus.FAIL);
    }

    @Test(expectedExceptions = AtlasBaseException.class)
    public void requestingExportOfNonExistentEntity_ReturnsFailure() throws Exception {
        AtlasExportRequest request = getRequestForEmployee();
        tamperEmployeeRequest(request);
        runExportWithParameters(request);
    }

    @Test
    public void requestForTypeFull() {
        AtlasExportRequest req = getRequestForTypeFull("Department,Employee");

        assertNotNull(req);
        assertEquals(req.getItemsToExport().size(), 1);
        assertEquals(req.getOptions().get(AtlasExportRequest.OPTION_ATTR_MATCH_TYPE), "forType");
    }

    @Test
    public void verifyTypeFull() throws AtlasBaseException, IOException {
        ZipSource zipSource = runExportWithParameters(getRequestForTypeFull("Department,Employee,Manager"));
        verifyExportForFullEmployeeData(zipSource);
    }

    private AtlasExportRequest getRequestForFullFetch() {
        AtlasExportRequest request = new AtlasExportRequest();

        List<AtlasObjectId> itemsToExport = new ArrayList<>();
        itemsToExport.add(new AtlasObjectId("hive_db", "qualifiedName", "default@cl1"));
        request.setItemsToExport(itemsToExport);

        return request;
    }

    private AtlasExportRequest getRequestForDept(boolean addFetchType, String fetchTypeValue, boolean addMatchType, String matchTypeValue) {
        AtlasExportRequest request = new AtlasExportRequest();

        List<AtlasObjectId> itemsToExport = new ArrayList<>();
        itemsToExport.add(new AtlasObjectId("Department", "name", "hr"));
        request.setItemsToExport(itemsToExport);

        setOptionsMap(request, addFetchType, fetchTypeValue, addMatchType, matchTypeValue);
        return request;
    }

    private AtlasExportRequest getRequestForEmployee() {
        AtlasExportRequest request = new AtlasExportRequest();

        List<AtlasObjectId> itemsToExport = new ArrayList<>();
        itemsToExport.add(new AtlasObjectId("Employee", "name", "Max"));
        request.setItemsToExport(itemsToExport);

        setOptionsMap(request, true, "CONNECTED", false, "");
        return request;
    }

    private void setOptionsMap(AtlasExportRequest request,
            boolean addFetchType, String fetchTypeValue, boolean addMatchType, String matchTypeValue) {
        Map<String, Object> optionsMap = null;
        if (addFetchType) {
            if (optionsMap == null) {
                optionsMap = new HashMap<>();
            }

            optionsMap.put("fetchType", fetchTypeValue);
            request.setOptions(optionsMap);
        }

        if (addMatchType) {
            if (optionsMap == null) {
                optionsMap = new HashMap<>();
            }

            optionsMap.put("matchType", matchTypeValue);
        }

        if (optionsMap != null) {
            request.setOptions(optionsMap);
        }
    }

    private ZipSource runExportWithParameters(AtlasExportRequest request) throws AtlasBaseException, IOException {
        final String requestingIP = "1.0.0.0";
        final String hostName     = "localhost";
        final String userName     = "admin";

        ByteArrayOutputStream baos    = new ByteArrayOutputStream();
        ZipSink               zipSink = new ZipSink(baos);
        AtlasExportResult     result  = exportService.run(zipSink, request, userName, hostName, requestingIP);

        zipSink.close();

        ByteArrayInputStream bis = new ByteArrayInputStream(baos.toByteArray());

        return new ZipSource(bis);
    }

    private AtlasExportRequest getRequestForTypeFull(String type) {
        String jsonRequest = "{ \"itemsToExport\": [ { \"typeName\": \"%s\" } ], \"options\": {  \"fetchType\": \"FULL\", \"matchType\": \"forType\"} }";
        return AtlasType.fromJson(String.format(jsonRequest, type), AtlasExportRequest.class);
    }

    private void tamperEmployeeRequest(AtlasExportRequest request) {
        AtlasObjectId objectId = request.getItemsToExport().get(0);
        objectId.getUniqueAttributes().remove("name");
        objectId.getUniqueAttributes().put("qualifiedName", "XXX@121");
    }

    private void verifyExportForEmployeeData(ZipSource zipSource) throws AtlasBaseException {
        final List<String> expectedEntityTypes = Arrays.asList("Manager", "Employee", "Department");

        assertNotNull(zipSource.getCreationOrder());
        assertEquals(zipSource.getCreationOrder().size(), 2);
        assertTrue(zipSource.hasNext());

        while (zipSource.hasNext()) {
            AtlasEntity entity = zipSource.next();

            assertNotNull(entity);
            assertEquals(entity.getStatus(), AtlasEntity.Status.ACTIVE);
            assertTrue(expectedEntityTypes.contains(entity.getTypeName()));
        }

        verifyTypeDefs(zipSource);
    }

    private void verifyExportForFullEmployeeData(ZipSource zipSource) throws AtlasBaseException {
        final List<String> expectedEntityTypes = Arrays.asList("Manager", "Employee", "Department");

        assertNotNull(zipSource.getCreationOrder());
        assertTrue(zipSource.hasNext());

        while (zipSource.hasNext()) {
            AtlasEntity entity = zipSource.next();

            assertNotNull(entity);
            assertEquals(entity.getStatus(), AtlasEntity.Status.ACTIVE);
            assertTrue(expectedEntityTypes.contains(entity.getTypeName()));
        }

        verifyTypeDefs(zipSource);
    }

    private void verifyExportForHrData(ZipSource zipSource) throws AtlasBaseException {
        assertNotNull(zipSource.getCreationOrder());
        assertEquals(zipSource.getCreationOrder().size(), 1);
        assertTrue(zipSource.hasNext());

        AtlasEntity entity = zipSource.next();

        assertNotNull(entity);
        assertEquals(entity.getTypeName(), "Department");
        assertEquals(entity.getStatus(), AtlasEntity.Status.ACTIVE);
        verifyTypeDefs(zipSource);
    }

    private void verifyExportForHrDataForConnected(ZipSource zipSource) throws AtlasBaseException {
        assertNotNull(zipSource.getCreationOrder());
        assertEquals(zipSource.getCreationOrder().size(), 1);
        assertTrue(zipSource.hasNext());

        AtlasEntity entity = zipSource.next();

        assertNotNull(entity);
        assertEquals(entity.getTypeName(), "Department");
        assertEquals(entity.getStatus(), AtlasEntity.Status.ACTIVE);
        verifyTypeDefs(zipSource);
    }

    private void verifyTypeDefs(ZipSource zipSource) throws AtlasBaseException {
        assertEquals(zipSource.getTypesDef().getEnumDefs().size(), 1);
        assertEquals(zipSource.getTypesDef().getClassificationDefs().size(), 0);
        assertEquals(zipSource.getTypesDef().getStructDefs().size(), 1);
        assertEquals(zipSource.getTypesDef().getEntityDefs().size(), 4);
    }
}
