/**
 * 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.store.graph.v2;

import org.apache.atlas.AtlasException;
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.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityExtInfo;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.instance.AtlasStruct;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
import org.apache.atlas.model.typedef.AtlasClassificationDef;
import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
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.graph.GraphBackedSearchIndexer;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
import org.apache.atlas.repository.store.graph.AtlasEntityStore;
import org.apache.atlas.repository.store.graph.v1.DeleteHandlerDelegate;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasArrayType;
import org.apache.atlas.type.AtlasMapType;
import org.apache.atlas.type.AtlasStructType;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.type.AtlasTypeUtil;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.Guice;

import javax.inject.Inject;

import java.util.Collections;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.TestUtilsV2.randomString;
import static org.mockito.Mockito.mock;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

@Guice(modules = TestModules.TestOnlyModule.class)
public class AtlasEntityTestBase extends AtlasTestBase {
    @Inject
    protected AtlasGraph graph;
    @Inject
    AtlasTypeRegistry typeRegistry;
    @Inject
    AtlasTypeDefStore typeDefStore;
    @Inject
    AtlasEntityStore entityStore;
    @Inject
    DeleteHandlerDelegate deleteDelegate;
    AtlasEntityChangeNotifier mockChangeNotifier = mock(AtlasEntityChangeNotifier.class);
    @Inject
    private EntityGraphMapper graphMapper;

    @BeforeClass
    public void setUp() throws Exception {
        RequestContext.clear();
        RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);

        super.initialize();

        new GraphBackedSearchIndexer(typeRegistry);
    }

    @AfterClass
    public void clear() throws Exception {
        AtlasGraphProvider.cleanup();

        super.cleanup();
    }

    @BeforeTest
    public void init() throws Exception {
        entityStore = new AtlasEntityStoreV2(graph, deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);

        RequestContext.clear();
        RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);
    }

    protected AtlasClassificationDef getTagWithName(AtlasTypesDef typesDef, String tagName, String attributeType) {
        AtlasClassificationDef aTag = new AtlasClassificationDef(tagName);
        AtlasAttributeDef attributeDef = new AtlasAttributeDef("testAttribute", attributeType, true,
                AtlasAttributeDef.Cardinality.SINGLE, 0, 1, false, true, false,
                Collections.emptyList());

        aTag.addAttribute(attributeDef);
        typesDef.setClassificationDefs(Collections.singletonList(aTag));
        return aTag;
    }

    protected void createTag(String tagName, String attributeType) {
        try {
            AtlasTypesDef typesDef = new AtlasTypesDef();
            getTagWithName(typesDef, tagName, attributeType);
            typeDefStore.createTypesDef(typesDef);
        } catch (AtlasBaseException e) {
            fail("Tag creation should've succeeded");
        }
    }

    protected String randomStrWithReservedChars() {
        return randomString() + "\"${}%";
    }

    protected void validateMutationResponse(EntityMutationResponse response, EntityOperation op, int expectedNumCreated) {
        List<AtlasEntityHeader> entitiesCreated = response.getEntitiesByOperation(op);
        assertNotNull(entitiesCreated);
        assertEquals(entitiesCreated.size(), expectedNumCreated);
    }

    protected void validateEntity(AtlasEntityExtInfo entityExtInfo, AtlasEntity actual) throws AtlasBaseException, AtlasException {
        validateEntity(entityExtInfo, actual, entityExtInfo.getEntity(actual.getGuid()));
    }

    protected void validateEntity(AtlasEntityExtInfo entityExtInfo, AtlasStruct actual, AtlasStruct expected) throws AtlasBaseException, AtlasException {
        if (expected == null) {
            assertNull(actual, "expected null instance. Found " + actual);

            return;
        }

        assertNotNull(actual, "found null instance");

        AtlasStructType entityType = (AtlasStructType) typeRegistry.getType(actual.getTypeName());
        for (String attrName : expected.getAttributes().keySet()) {
            Object expectedVal = expected.getAttribute(attrName);
            Object actualVal   = actual.getAttribute(attrName);

            AtlasType attrType = entityType.getAttributeType(attrName);
            validateAttribute(entityExtInfo, actualVal, expectedVal, attrType, attrName);
        }
    }

    protected void validateAttribute(AtlasEntityExtInfo entityExtInfo, Object actual, Object expected, AtlasType attributeType, String attrName) throws AtlasBaseException, AtlasException {
        switch (attributeType.getTypeCategory()) {
            case OBJECT_ID_TYPE:
                assertTrue(actual instanceof AtlasObjectId);
                String guid = ((AtlasObjectId) actual).getGuid();
                assertTrue(AtlasTypeUtil.isAssignedGuid(guid), "expected assigned guid. found " + guid);
                break;

            case PRIMITIVE:
            case ENUM:
                assertEquals(actual, expected);
                break;

            case MAP:
                AtlasMapType mapType = (AtlasMapType) attributeType;
                AtlasType valueType = mapType.getValueType();
                Map actualMap = (Map) actual;
                Map expectedMap = (Map) expected;

                if (MapUtils.isNotEmpty(expectedMap)) {
                    assertTrue(MapUtils.isNotEmpty(actualMap));

                    // deleted entries are included in the attribute; hence use >=
                    assertTrue(actualMap.size() >= expectedMap.size());

                    for (Object key : expectedMap.keySet()) {
                        validateAttribute(entityExtInfo, actualMap.get(key), expectedMap.get(key), valueType, attrName);
                    }
                }
                break;

            case ARRAY:
                AtlasArrayType arrType = (AtlasArrayType) attributeType;
                AtlasType elemType = arrType.getElementType();
                List actualList = (List) actual;
                List expectedList = (List) expected;

                if (CollectionUtils.isNotEmpty(expectedList)) {
                    assertTrue(CollectionUtils.isNotEmpty(actualList));

                    //actual list could have deleted entities. Hence size may not match.
                    assertTrue(actualList.size() >= expectedList.size());

                    for (Object o : expectedList) {
                        assertTrue(actualList.contains(o));
                    }
                }
                break;
            case STRUCT:
                AtlasStruct expectedStruct = (AtlasStruct) expected;
                AtlasStruct actualStruct = (AtlasStruct) actual;

                validateEntity(entityExtInfo, actualStruct, expectedStruct);
                break;
            default:
                fail("Unknown type category");
        }
    }

    protected AtlasEntity getEntityFromStore(AtlasEntityHeader header) throws AtlasBaseException {
        return header != null ? getEntityFromStore(header.getGuid()) : null;
    }

    protected AtlasEntity getEntityFromStore(String guid) throws AtlasBaseException {
        AtlasEntityWithExtInfo entity = guid != null ? entityStore.getById(guid) : null;

        return entity != null ? entity.getEntity() : null;
    }

    protected void createTypesDef(AtlasTypesDef[] testTypesDefs) throws AtlasBaseException {
        for (AtlasTypesDef typesDef : testTypesDefs) {
            AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);

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