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

import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.instance.AtlasRelatedObjectId;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.store.graph.v2.EntityStream;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasTypeRegistry;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public final class EntityGraphDiscoveryContext {
    private final AtlasTypeRegistry               typeRegistry;
    private final EntityStream                    entityStream;
    private final List<String>                    referencedGuids          = new ArrayList<>();
    private final Set<AtlasObjectId>              referencedByUniqAttribs  = new HashSet<>();
    private final Map<String, AtlasVertex>        resolvedGuids            = new HashMap<>();
    private final Map<AtlasObjectId, AtlasVertex> resolvedIdsByUniqAttribs = new HashMap<>();
    private final Set<String>                     localGuids               = new HashSet<>();

    public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
        this.typeRegistry = typeRegistry;
        this.entityStream = entityStream;
    }

    public EntityStream getEntityStream() {
        return entityStream;
    }

    public List<String> getReferencedGuids() {
        return referencedGuids;
    }

    public Set<AtlasObjectId> getReferencedByUniqAttribs() {
        return referencedByUniqAttribs;
    }

    public Map<String, AtlasVertex> getResolvedGuids() {
        return resolvedGuids;
    }

    public Map<AtlasObjectId, AtlasVertex> getResolvedIdsByUniqAttribs() {
        return resolvedIdsByUniqAttribs;
    }

    public Set<String> getLocalGuids() {
        return localGuids;
    }

    public void addReferencedGuid(String guid) {
        if (!referencedGuids.contains(guid)) {
            referencedGuids.add(guid);
        }
    }

    public void addReferencedByUniqAttribs(AtlasObjectId objId) {
        referencedByUniqAttribs.add(objId);
    }

    public void addResolvedGuid(String guid, AtlasVertex vertex) {
        resolvedGuids.put(guid, vertex);
    }

    public void addResolvedIdByUniqAttribs(AtlasObjectId objId, AtlasVertex vertex) {
        resolvedIdsByUniqAttribs.put(objId, vertex);
    }

    public void addLocalGuidReference(String guid) {
        localGuids.add(guid);
    }

    public boolean isResolvedGuid(String guid) {
        return resolvedGuids.containsKey(guid);
    }

    public boolean isResolvedIdByUniqAttrib(AtlasObjectId objId) {
        return resolvedIdsByUniqAttribs.containsKey(objId);
    }

    public AtlasVertex getResolvedEntityVertex(String guid) {
        return resolvedGuids.get(guid);
    }

    public AtlasVertex getResolvedEntityVertex(AtlasObjectId objId) {
        if (resolvedIdsByUniqAttribs.containsKey(objId)) {
            return getAtlasVertexFromResolvedIdsByAttribs(objId);
        } else if (objId instanceof AtlasRelatedObjectId) {
            objId = new AtlasObjectId(objId.getGuid(), objId.getTypeName(), objId.getUniqueAttributes());
        }

        return getAtlasVertexFromResolvedIdsByAttribs(objId);
    }

    public StringBuilder toString(StringBuilder sb) {
        if (sb == null) {
            sb = new StringBuilder();
        }

        sb.append("EntityGraphDiscoveryCtx{");
        sb.append("referencedGuids=").append(referencedGuids);
        sb.append(", referencedByUniqAttribs=").append(referencedByUniqAttribs);
        sb.append(", resolvedGuids='").append(resolvedGuids);
        sb.append(", resolvedIdsByUniqAttribs='").append(resolvedIdsByUniqAttribs);
        sb.append(", localGuids='").append(localGuids);
        sb.append('}');

        return sb;
    }

    @Override
    public String toString() {
        return toString(new StringBuilder()).toString();
    }

    public void cleanUp() {
        referencedGuids.clear();
        referencedByUniqAttribs.clear();
        resolvedGuids.clear();
        resolvedIdsByUniqAttribs.clear();
        localGuids.clear();
    }

    private AtlasVertex getAtlasVertexFromResolvedIdsByAttribs(AtlasObjectId objId) {
        AtlasVertex vertex = resolvedIdsByUniqAttribs.get(objId);

        // check also for sub-types; ref={typeName=Asset; guid=abcd} should match {typeName=hive_table; guid=abcd}
        if (vertex == null) {
            final AtlasEntityType entityType  = typeRegistry.getEntityTypeByName(objId.getTypeName());
            final Set<String>     allSubTypes = entityType.getAllSubTypes();

            for (String subType : allSubTypes) {
                AtlasObjectId subTypeObjId = new AtlasObjectId(objId.getGuid(), subType, objId.getUniqueAttributes());

                vertex = resolvedIdsByUniqAttribs.get(subTypeObjId);

                if (vertex != null) {
                    resolvedIdsByUniqAttribs.put(objId, vertex);
                    break;
                }
            }
        }

        return vertex;
    }
}
