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

import org.apache.atlas.DeleteType;
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.AtlasObjectId;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.typedef.AtlasEntityDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.graphdb.AtlasEdge;
import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.store.graph.AtlasEntityStore;
import org.apache.atlas.repository.store.graph.v2.AtlasEntityStream;
import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.utils.TestResourceFileUtils;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Guice;
import org.testng.annotations.Test;

import javax.inject.Inject;

import java.io.IOException;
import java.util.Iterator;
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.assertTrue;

@Guice(modules = TestModules.TestOnlyModule.class)
public class SoftReferenceTest {
    private static final String TYPE_RDBMS_DB             = "rdbms_db";
    private static final String RDBMS_DB_FILE             = "rdbms-db";
    private static final String TYPE_RDBMS_STORAGE        = "rdbms_storage";
    private static final String TYPESDEF_FILE_NAME        = "typesDef-soft-ref";
    private static final String RDBMS_DB_STORAGE_PROPERTY = "sd";
    private static final String RDBMS_DB_TABLES_PROPERTY  = "tables";
    private static final String RDBMS_DB_REGIONS_PROPERTY = "regions";
    private static final String RDBMS_SD_PROPERTY         = "rdbms_db.sd";
    private static final String TYPE_RDBMS_TABLES         = "rdbms_table";

    @Inject
    AtlasTypeRegistry typeRegistry;

    @Inject
    private AtlasTypeDefStore typeDefStore;

    @Inject
    private AtlasEntityStore entityStore;

    private AtlasType dbType;
    private String    dbGuid;
    private String    storageGuid;

    @BeforeMethod
    public void init() throws Exception {
        RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);
        RequestContext.get().setDeleteType(DeleteType.SOFT);
    }

    @Test
    public void typeCreationFromFile() throws IOException, AtlasBaseException {
        String        typesDefJson = TestResourceFileUtils.getJson(TYPESDEF_FILE_NAME);
        AtlasTypesDef typesDef     = AtlasType.fromJson(typesDefJson, AtlasTypesDef.class);

        assertNotNull(typesDef);

        typeDefStore.createTypesDef(typesDef);

        dbType = typeRegistry.getType(TYPE_RDBMS_DB);

        assertNotNull(dbType);

        AtlasEntityDef dbType = typeRegistry.getEntityDefByName(TYPE_RDBMS_DB);

        assertNotNull(dbType);
        assertTrue(dbType.getAttribute(RDBMS_DB_STORAGE_PROPERTY).isSoftReferenced());
        assertTrue(dbType.getAttribute(RDBMS_DB_TABLES_PROPERTY).isSoftReferenced());
        assertTrue(dbType.getAttribute(RDBMS_DB_REGIONS_PROPERTY).isSoftReferenced());
        assertNotNull(typeRegistry.getEntityDefByName(TYPE_RDBMS_STORAGE));
        assertNotNull(typeRegistry.getEntityDefByName(TYPE_RDBMS_TABLES));
    }

    @Test(dependsOnMethods = "typeCreationFromFile")
    public void entityCreationUsingSoftRef() throws IOException, AtlasBaseException {
        final int                          expectedEntityCount = 6;
        AtlasEntity.AtlasEntityWithExtInfo dbEntity            = AtlasType.fromJson(TestResourceFileUtils.getJson(RDBMS_DB_FILE), AtlasEntity.AtlasEntityWithExtInfo.class);
        EntityMutationResponse             response            = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity), false);

        assertNotNull(response);
        assertEquals(response.getCreatedEntities().size(), expectedEntityCount);
        assertGraphStructure(response.getCreatedEntities().get(0).getGuid(), response.getCreatedEntities().get(1).getGuid(), RDBMS_SD_PROPERTY);

        dbGuid      = response.getCreatedEntities().get(0).getGuid();
        storageGuid = response.getCreatedEntities().get(1).getGuid();
    }

    @Test(dependsOnMethods = "entityCreationUsingSoftRef")
    public void deletetingCollections() throws AtlasBaseException {
        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(dbGuid);

        assertNotNull(entityWithExtInfo);

        List list = (List) entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_TABLES_PROPERTY);

        list.remove(1);

        Map map = (Map) entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_REGIONS_PROPERTY);

        map.remove("east");

        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo), true);

        assertNotNull(response);
        assertTrue(response.getPartialUpdatedEntities().size() > 0);
        assertAttribute(dbGuid, storageGuid, 1, 1);
    }

    @Test(dependsOnMethods = "deletetingCollections")
    public void addingCollections() throws AtlasBaseException {
        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(dbGuid);

        assertNotNull(entityWithExtInfo);
        addNewTables(entityWithExtInfo);
        addNewRegions(entityWithExtInfo);

        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo), true);

        assertNotNull(response);
        assertTrue(response.getPartialUpdatedEntities().size() > 0);
        assertAttribute(dbGuid, storageGuid, 3, 3);
    }

    private void addNewRegions(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
        Map map = (Map) entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_REGIONS_PROPERTY);

        AtlasEntity region1 = getDefaultTableEntity("r1");
        AtlasEntity region2 = getDefaultTableEntity("r2");

        map.put("north", new AtlasObjectId(region1.getGuid(), region1.getTypeName()));
        map.put("south", new AtlasObjectId(region2.getGuid(), region2.getTypeName()));

        entityWithExtInfo.addReferredEntity(region1);
        entityWithExtInfo.addReferredEntity(region2);
    }

    private void addNewTables(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
        List        list   = (List) entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_TABLES_PROPERTY);
        AtlasEntity table1 = getDefaultTableEntity("newTable-1");
        AtlasEntity table2 = getDefaultTableEntity("newTable-2");

        entityWithExtInfo.addReferredEntity(table1);
        entityWithExtInfo.addReferredEntity(table2);

        list.add(new AtlasObjectId(table1.getGuid(), table1.getTypeName()));
        list.add(new AtlasObjectId(table2.getGuid(), table2.getTypeName()));
    }

    private AtlasEntity getDefaultTableEntity(String name) throws AtlasBaseException {
        AtlasEntityType type = (AtlasEntityType) typeRegistry.getType(TYPE_RDBMS_TABLES);
        AtlasEntity     ret  = type.createDefaultValue();

        ret.setAttribute("name", name);

        return ret;
    }

    private void assertGraphStructure(String dbGuid, String storageGuid, String propertyName) throws AtlasBaseException {
        AtlasVertex         vertex   = AtlasGraphUtilsV2.findByGuid(dbGuid);
        Iterator<AtlasEdge> edgesOut = vertex.getEdges(AtlasEdgeDirection.OUT).iterator();
        Iterator<AtlasEdge> edgesIn  = vertex.getEdges(AtlasEdgeDirection.IN).iterator();
        String              sd       = AtlasGraphUtilsV2.getProperty(vertex, propertyName, String.class);

        assertNotNull(sd);
        assertAttribute(dbGuid, storageGuid, 2, 2);
        assertFalse(edgesOut.hasNext());
        assertFalse(edgesIn.hasNext());
        assertNotNull(vertex);
    }

    private void assertAttribute(String dbGuid, String storageGuid, int expectedTableCount, int expectedRegionCount) throws AtlasBaseException {
        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(dbGuid);
        AtlasEntity                        entity            = entityWithExtInfo.getEntity();
        Object                             val               = entity.getAttribute(RDBMS_DB_STORAGE_PROPERTY);

        assertTrue(val instanceof AtlasObjectId);
        assertEquals(((AtlasObjectId) val).getTypeName(), TYPE_RDBMS_STORAGE);
        assertEquals(((AtlasObjectId) val).getGuid(), storageGuid);
        assertNotNull(entity.getAttribute(RDBMS_DB_TABLES_PROPERTY));
        assertEquals(((List) entity.getAttribute(RDBMS_DB_TABLES_PROPERTY)).size(), expectedTableCount);
        assertNotNull(entity.getAttribute(RDBMS_DB_REGIONS_PROPERTY));
        assertEquals(((Map) entity.getAttribute(RDBMS_DB_REGIONS_PROPERTY)).size(), expectedRegionCount);
    }
}
