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

import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.instance.AtlasClassification;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.AtlasStruct;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.instance.EntityMutations;
import org.apache.atlas.type.AtlasTypeUtil;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static java.util.Arrays.asList;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_COLUMNS;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_COMMENT;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_COMPRESSED;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_CREATE_TIME;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_DATA_TYPE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_DB;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_END_TIME;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_INPUTS;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_LAST_ACCESS_TIME;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_LEVEL;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_LOCATION_URI;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_OPERATION_TYPE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_OUTPUTS;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_OWNER;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_GRAPH;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_ID;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_PLAN;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_TEXT;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE1;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE2;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_START_TIME;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_TABLE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_TABLE_TYPE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_USERNAME;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.STRUCT_TYPE_SERDE;
import static org.apache.atlas.examples.sampleapp.SampleAppConstants.TABLE_DATABASE_TYPE;
import static org.apache.atlas.type.AtlasTypeUtil.toAtlasRelatedObjectId;
import static org.apache.atlas.type.AtlasTypeUtil.toAtlasRelatedObjectIds;

public class EntityExample {
    private static final String DATABASE_NAME                = "employee_db_entity";
    private static final String TABLE_NAME                   = "employee_table_entity";
    private static final String PROCESS_NAME                 = "employee_process_entity";
    private static final String METADATA_NAMESPACE_SUFFIX    = "@cl1";
    private static final String MANAGED_TABLE                = "Managed";
    private static final String ATTR_NAME                    = "name";
    private static final String ATTR_DESCRIPTION             = "description";
    private static final String ATTR_QUALIFIED_NAME          = "qualifiedName";
    private static final String REFERENCEABLE_ATTRIBUTE_NAME = ATTR_QUALIFIED_NAME;
    private static final String COLUMN_TIME_ID               = "time_id";
    private static final String COLUMN_CUSTOMER_ID           = "customer_id";
    private static final String COLUMN_COMPANY_ID            = "company_id";

    private final AtlasClientV2     client;
    private       AtlasEntity       dbEntity;
    private       AtlasEntity       tableEntityUS;
    private       AtlasEntity       tableEntityCanada;
    private       AtlasEntityHeader loadProcess;

    EntityExample(AtlasClientV2 client) {
        this.client = client;
    }

    public void createEntities() throws Exception {
        if (dbEntity == null) {
            dbEntity = createDatabaseEntity(DATABASE_NAME);

            SampleApp.log("Created entity: typeName=" + dbEntity.getTypeName() + ", qualifiedName=" + dbEntity.getAttribute(ATTR_QUALIFIED_NAME) + ", guid=" + dbEntity.getGuid());
        }

        if (tableEntityCanada == null) {
            tableEntityCanada = createTableEntity(TABLE_NAME + "_CANADA");

            SampleApp.log("Created entity: typeName=" + tableEntityCanada.getTypeName() + ", qualifiedName=" + tableEntityCanada.getAttribute(ATTR_QUALIFIED_NAME) + ", guid=" + tableEntityCanada.getGuid());
        }

        if (tableEntityUS == null) {
            tableEntityUS = createTableEntity(TABLE_NAME + "_US");

            SampleApp.log("Created entity: typeName=" + tableEntityUS.getTypeName() + ", qualifiedName=" + tableEntityUS.getAttribute(ATTR_QUALIFIED_NAME) + ", guid=" + tableEntityUS.getGuid());
        }

        if (loadProcess == null) {
            loadProcess = createProcessEntity(PROCESS_NAME);

            SampleApp.log("Created entity: typeName=" + loadProcess.getTypeName() + ", qualifiedName=" + loadProcess.getAttribute(ATTR_QUALIFIED_NAME) + ", guid=" + loadProcess.getGuid());
        }
    }

    public AtlasEntity getTableEntity() {
        return tableEntityUS;
    }

    public void getEntityByGuid(String entityGuid) throws Exception {
        AtlasEntityWithExtInfo entity = client.getEntityByGuid(entityGuid);

        if (entity != null) {
            SampleApp.log("Retrieved entity with guid=" + entityGuid);
            SampleApp.log("  " + entity);
        }
    }

    public void deleteEntities() throws Exception {
        client.deleteEntityByGuid(loadProcess.getGuid());

        SampleApp.log("Deleted entity: guid=" + loadProcess.getGuid());

        List<String> entityGuids = Arrays.asList(tableEntityUS.getGuid(), tableEntityCanada.getGuid(), dbEntity.getGuid());

        client.deleteEntitiesByGuids(entityGuids);

        SampleApp.log("Deleted entities:");
        for (String entityGuid : entityGuids) {
            SampleApp.log("  guid=" + entityGuid);
        }
    }

    protected AtlasEntity createHiveDBInstance(String dbName) {
        AtlasEntity entity = new AtlasEntity(SampleAppConstants.DATABASE_TYPE);

        entity.setAttribute(ATTR_NAME, dbName);
        entity.setAttribute(ATTR_DESCRIPTION, "employee database");
        entity.setAttribute(METADATA_NAMESPACE_SUFFIX, "employeeCluster");
        entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, dbName + METADATA_NAMESPACE_SUFFIX);
        entity.setAttribute(ATTR_OWNER, "user");
        entity.setAttribute(ATTR_LOCATION_URI, "/tmp");
        entity.setAttribute(ATTR_CREATE_TIME, 1000);

        return entity;
    }

    private AtlasEntity createTableEntity(String tableName) throws Exception {
        return createHiveTable(dbEntity, tableName, MANAGED_TABLE,
                Arrays.asList(createColumn(COLUMN_TIME_ID, "int", "time id"),
                        createColumn(COLUMN_CUSTOMER_ID, "int", "customer id", SampleAppConstants.PII_TAG),
                        createColumn(COLUMN_COMPANY_ID, "double", "company id", SampleAppConstants.FINANCE_TAG)),
                SampleAppConstants.METRIC_TAG);
    }

    private AtlasEntityHeader createProcessEntity(String processName) throws Exception {
        return createProcess(processName, "hive query for monthly avg salary", "user ETL",
                asList(tableEntityUS),
                asList(tableEntityCanada),
                "create table as select ", "plan", "id", "graph", SampleAppConstants.CLASSIFIED_TAG);
    }

    private AtlasEntityHeader createProcess(String name, String description, String user, List<AtlasEntity> inputs, List<AtlasEntity> outputs,
            String queryText, String queryPlan, String queryId, String queryGraph, String... classificationNames) throws Exception {
        AtlasEntity entity = new AtlasEntity(SampleAppConstants.PROCESS_TYPE);

        entity.setAttribute(ATTR_NAME, name);
        entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + METADATA_NAMESPACE_SUFFIX);
        entity.setAttribute(ATTR_DESCRIPTION, description);
        entity.setAttribute(ATTR_USERNAME, user);
        entity.setAttribute(ATTR_START_TIME, System.currentTimeMillis());
        entity.setAttribute(ATTR_END_TIME, System.currentTimeMillis() + 10000);
        entity.setAttribute(ATTR_QUERY_TEXT, queryText);
        entity.setAttribute(ATTR_QUERY_PLAN, queryPlan);
        entity.setAttribute(ATTR_QUERY_ID, queryId);
        entity.setAttribute(ATTR_QUERY_GRAPH, queryGraph);
        entity.setAttribute(ATTR_OPERATION_TYPE, "testOperation");

        entity.setRelationshipAttribute(ATTR_INPUTS, toAtlasRelatedObjectIds(inputs));
        entity.setRelationshipAttribute(ATTR_OUTPUTS, toAtlasRelatedObjectIds(outputs));

        entity.setClassifications(toAtlasClassifications(classificationNames));

        return createEntity(new AtlasEntityWithExtInfo(entity));
    }

    private AtlasEntity createColumn(String name, String dataType, String comment, String... classificationNames) {
        AtlasEntity ret = new AtlasEntity(SampleAppConstants.COLUMN_TYPE);

        ret.setAttribute(ATTR_NAME, name);
        ret.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + METADATA_NAMESPACE_SUFFIX);
        ret.setAttribute(ATTR_DATA_TYPE, dataType);
        ret.setAttribute(ATTR_COMMENT, comment);

        ret.setClassifications(toAtlasClassifications(classificationNames));

        return ret;
    }

    private List<AtlasClassification> toAtlasClassifications(String[] classificationNames) {
        List<AtlasClassification> ret = new ArrayList<>();

        if (classificationNames != null) {
            for (String classificationName : classificationNames) {
                ret.add(new AtlasClassification(classificationName));
            }
        }

        return ret;
    }

    private AtlasEntityHeader createEntity(AtlasEntityWithExtInfo atlasEntityWithExtInfo) {
        EntityMutationResponse entity;

        try {
            entity = client.createEntity(atlasEntityWithExtInfo);

            if (entity != null && entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE) != null) {
                List<AtlasEntityHeader> list = entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

                if (list.size() > 0) {
                    return entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0);
                }
            }
        } catch (AtlasServiceException e) {
            SampleApp.log("failed in create entity");
            e.printStackTrace();
        }

        return null;
    }

    private AtlasEntity createDatabaseEntity(String dbName) {
        AtlasEntity       hiveDBInstance = createHiveDBInstance(dbName);
        AtlasEntityHeader entityHeader   = createEntity(new AtlasEntityWithExtInfo(hiveDBInstance));

        if (entityHeader != null && entityHeader.getGuid() != null) {
            hiveDBInstance.setGuid(entityHeader.getGuid());
        }

        return hiveDBInstance;
    }

    private AtlasEntity createHiveTable(AtlasEntity database, String tableName, String tableType, List<AtlasEntity> columns, String... classificationNames) throws Exception {
        AtlasEntityWithExtInfo entityWithExtInfo = new AtlasEntityWithExtInfo();
        AtlasEntity            hiveTableInstance = createHiveTable(database, tableName, tableType, classificationNames);

        entityWithExtInfo.setEntity(hiveTableInstance);
        hiveTableInstance.setRelationshipAttribute(ATTR_COLUMNS, toAtlasRelatedObjectIds(columns));

        for (AtlasEntity column : columns) {
            column.setRelationshipAttribute(ATTR_TABLE, toAtlasRelatedObjectId(hiveTableInstance));
            entityWithExtInfo.addReferredEntity(column);
        }

        AtlasEntityHeader createdHeader = createEntity(entityWithExtInfo);

        if (createdHeader != null && createdHeader.getGuid() != null) {
            hiveTableInstance.setGuid(createdHeader.getGuid());
        }

        return hiveTableInstance;
    }

    private AtlasEntity createHiveTable(AtlasEntity database, String tableName, String tableType, String... classificationNames) throws Exception {
        AtlasEntity table = new AtlasEntity(SampleAppConstants.TABLE_TYPE);

        table.setAttribute(ATTR_NAME, tableName);
        table.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, database.getAttribute(ATTR_NAME) + "." + tableName + METADATA_NAMESPACE_SUFFIX);
        table.setAttribute(ATTR_TABLE_TYPE, tableType);
        table.setRelationshipAttribute(ATTR_DB, AtlasTypeUtil.getAtlasRelatedObjectId(database, TABLE_DATABASE_TYPE));

        table.setAttribute(ATTR_DESCRIPTION, "emp table");
        table.setAttribute(ATTR_LAST_ACCESS_TIME, "2014-07-11T08:00:00.000Z");
        table.setAttribute(ATTR_LEVEL, 2);
        table.setAttribute(ATTR_COMPRESSED, false);
        table.setClassifications(toAtlasClassifications(classificationNames));

        AtlasStruct serde1 = new AtlasStruct(STRUCT_TYPE_SERDE);

        serde1.setAttribute(ATTR_NAME, "serde1");
        serde1.setAttribute(ATTR_SERDE, "serde1");
        table.setAttribute(ATTR_SERDE1, serde1);

        AtlasStruct serde2 = new AtlasStruct(STRUCT_TYPE_SERDE);

        serde2.setAttribute(ATTR_NAME, "serde2");
        serde2.setAttribute(ATTR_SERDE, "serde2");
        table.setAttribute(ATTR_SERDE2, serde2);

        return table;
    }
}
