/**
 * 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.tagpropagation;

import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.RequestContext;
import org.apache.atlas.TestModules;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.exception.EntityNotFoundException;
import org.apache.atlas.model.instance.AtlasClassification;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.typedef.AtlasClassificationDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.AtlasTestBase;
import org.apache.atlas.repository.graph.GraphHelper;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.impexp.ImportService;
import org.apache.atlas.repository.impexp.ZipFileResourceTestUtils;
import org.apache.atlas.repository.store.graph.AtlasEntityStore;
import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
import org.apache.atlas.repository.store.graph.v2.EntityGraphMapper;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.tasks.TaskManagement;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.commons.lang3.StringUtils;
import org.testng.SkipException;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Guice;
import org.testng.annotations.Test;

import javax.inject.Inject;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertex;
import static org.apache.atlas.repository.graph.GraphHelper.getPropagatedTraitNames;
import static org.apache.atlas.repository.impexp.ZipFileResourceTestUtils.runImportWithNoParameters;
import static org.apache.atlas.utils.TestLoadModelUtils.loadModelFromJson;
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 ClassificationPropagationWithTasksTest extends AtlasTestBase {
    private static final String IMPORT_FILE = "tag-propagation-data.zip";

    private static final String IMPORT_DELETE_FILE = "deleted_tab_propagation.zip";

    private static final String HDFS_PATH_EMPLOYEES = "a3955120-ac17-426f-a4af-972ec8690e5f";

    private static final String HIVE_TABLE = "089c1ad4-9dde-4f9e-80c8-12a3046be337";

    private static final String HIVE_TABLE_CTAS = "e83551b7-bbef-45aa-99d5-3d98c0ac737b";

    @Inject
    private AtlasTypeDefStore typeDefStore;

    @Inject
    private AtlasTypeRegistry typeRegistry;

    @Inject
    private AtlasEntityStore entityStore;

    @Inject
    private ImportService importService;

    @Inject
    private EntityGraphMapper entityGraphMapper;

    @Inject
    private TaskManagement tasksManagement;

    public static InputStream getZipSource(String fileName) throws IOException {
        return ZipFileResourceTestUtils.getFileInputStream(fileName);
    }

    @BeforeClass
    public void setup() throws Exception {
        RequestContext.clear();

        super.initialize();

        this.tasksManagement.start();
        entityGraphMapper.setTasksUseFlag(true);

        loadModelFilesAndImportTestData();
    }

    @Test
    public void parameterValidation() throws AtlasBaseException {
        try {
            entityGraphMapper.propagateClassification(null, null, null);
            entityGraphMapper.propagateClassification("unknown", "abcd", "xyz");
        } catch (AtlasBaseException e) {
            assertNotNull(e.getCause());
            assertTrue(e.getCause() instanceof EntityNotFoundException);
        }

        List<String> ret = entityGraphMapper.propagateClassification(HDFS_PATH_EMPLOYEES, StringUtils.EMPTY, StringUtils.EMPTY);

        assertNull(ret);

        ret = entityGraphMapper.deleteClassificationPropagation(StringUtils.EMPTY, StringUtils.EMPTY);

        assertNull(ret);

        AtlasEntity hdfsEmployees = getEntity(HDFS_PATH_EMPLOYEES);

        ret = entityGraphMapper.propagateClassification(hdfsEmployees.getGuid(), StringUtils.EMPTY, StringUtils.EMPTY);

        assertNull(ret);
    }

    @Test
    public void add() throws AtlasBaseException {
        final String tagNameX = "tagX";
        final String tagNameY = "tagY";

        AtlasEntity hdfsEmployees = getEntity(HDFS_PATH_EMPLOYEES);

        AtlasClassification tagX = new AtlasClassification(tagNameX);

        tagX.setEntityGuid(hdfsEmployees.getGuid());
        tagX.setPropagate(true);

        AtlasClassification tagY = new AtlasClassification(tagNameY);

        tagY.setEntityGuid(hdfsEmployees.getGuid());
        tagY.setPropagate(false);

        entityStore.addClassification(Collections.singletonList(HDFS_PATH_EMPLOYEES), tagX);
        entityStore.addClassification(Collections.singletonList(HDFS_PATH_EMPLOYEES), tagY);

        AtlasVertex entityVertex         = AtlasGraphUtilsV2.findByGuid(hdfsEmployees.getGuid());
        AtlasVertex classificationVertex = GraphHelper.getClassificationVertex(entityVertex, tagNameX);

        assertNotNull(entityVertex);
        assertNotNull(classificationVertex);

        AtlasEntity entityUpdated = getEntity(HDFS_PATH_EMPLOYEES);

        assertNotNull(entityUpdated.getPendingTasks());

        List<String> impactedEntities = entityGraphMapper.propagateClassification(hdfsEmployees.getGuid(), classificationVertex.getId().toString(), StringUtils.EMPTY);

        assertNotNull(impactedEntities);
    }

    @Test(dependsOnMethods = "add")
    public void update() throws AtlasBaseException {
        final String tagNameY = "tagY";

        AtlasEntity         hdfsEmployees = getEntity(HDFS_PATH_EMPLOYEES);
        AtlasClassification tagY          = new AtlasClassification(tagNameY);

        tagY.setEntityGuid(hdfsEmployees.getGuid());
        tagY.setPropagate(true);

        entityStore.updateClassifications(hdfsEmployees.getGuid(), Collections.singletonList(tagY));

        AtlasVertex entityVertex         = AtlasGraphUtilsV2.findByGuid(hdfsEmployees.getGuid());
        AtlasVertex classificationVertex = GraphHelper.getClassificationVertex(entityVertex, tagNameY);

        assertNotNull(RequestContext.get().getQueuedTasks());
        assertFalse(RequestContext.get().getQueuedTasks().isEmpty(), "No tasks were queued!");

        assertNotNull(entityVertex);
        assertNotNull(classificationVertex);
    }

    @Test(dependsOnMethods = "update")
    public void delete() throws AtlasBaseException {
        final String tagName = "tagX";

        AtlasEntity hdfsEmployees = getEntity(HDFS_PATH_EMPLOYEES);

        entityGraphMapper.propagateClassification(hdfsEmployees.getGuid(), StringUtils.EMPTY, StringUtils.EMPTY);

        AtlasClassification tagX = new AtlasClassification(tagName);

        tagX.setEntityGuid(hdfsEmployees.getGuid());
        tagX.setPropagate(false);

        AtlasVertex entityVertex         = AtlasGraphUtilsV2.findByGuid(hdfsEmployees.getGuid());
        AtlasVertex classificationVertex = GraphHelper.getClassificationVertex(entityVertex, tagName);

        try {
            entityStore.deleteClassification(HDFS_PATH_EMPLOYEES, tagX.getTypeName());
        } catch (AtlasBaseException e) {
            assertEquals(e.getAtlasErrorCode(), AtlasErrorCode.DELETE_TAG_PROPAGATION_NOT_ALLOWED);
        }

        assertNotNull(entityVertex);
        assertNotNull(classificationVertex);

        List<String> impactedEntities = entityGraphMapper.deleteClassificationPropagation(hdfsEmployees.getGuid(), classificationVertex.getId().toString());

        assertNotNull(impactedEntities);
    }

    @Test
    public void runImportForDeletedEntityLineage() throws Exception {
        runImportWithNoParameters(importService, getZipSource(IMPORT_DELETE_FILE));
        final String tagName = "classification1";

        AtlasEntity         hiveTable = getEntity(HIVE_TABLE);
        AtlasEntity         hiveTableCtas = getEntity(HIVE_TABLE_CTAS);

        AtlasVertex parentEntityVertex   = AtlasGraphUtilsV2.findByGuid(hiveTable.getGuid());

        AtlasVertex entityVertex         = AtlasGraphUtilsV2.findByGuid(hiveTableCtas.getGuid());

        AtlasVertex classificationVertex = getClassificationVertex(parentEntityVertex, tagName);
        assertNotNull(entityVertex);
        assertNotNull(parentEntityVertex);
        assertNotNull(classificationVertex);

        List<String> propagatedTraitNames = getPropagatedTraitNames(entityVertex);

        assertNotNull(propagatedTraitNames);
    }

    private void loadModelFilesAndImportTestData() {
        try {
            loadModelFromJson("0000-Area0/0010-base_model.json", typeDefStore, typeRegistry);
            loadModelFromJson("1000-Hadoop/1020-fs_model.json", typeDefStore, typeRegistry);
            loadModelFromJson("1000-Hadoop/1030-hive_model.json", typeDefStore, typeRegistry);

            loadSampleClassificationDefs();

            runImportWithNoParameters(importService, getZipSource(IMPORT_FILE));
        } catch (AtlasBaseException | IOException e) {
            throw new SkipException("Model loading failed!");
        }
    }

    private void loadSampleClassificationDefs() throws AtlasBaseException {
        AtlasClassificationDef tagX = new AtlasClassificationDef("tagX");
        AtlasClassificationDef tagY = new AtlasClassificationDef("tagY");

        typeDefStore.createTypesDef(new AtlasTypesDef(Collections.emptyList(), Collections.emptyList(), Arrays.asList(tagX, tagY), Collections.emptyList(), Collections.emptyList()));
    }

    private AtlasEntity getEntity(String entityGuid) throws AtlasBaseException {
        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(entityGuid);

        return entityWithExtInfo.getEntity();
    }
}
