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

import com.sun.jersey.api.client.ClientResponse;
import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.discovery.AtlasSearchResult;
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.EntityMutationResponse;
import org.apache.atlas.model.instance.EntityMutations;
import org.apache.atlas.trino.model.Catalog;
import org.apache.atlas.type.AtlasTypeUtil;
import org.apache.atlas.utils.AuthenticationUtil;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.apache.atlas.type.AtlasTypeUtil.ATTRIBUTE_QUALIFIED_NAME;

public class AtlasClientHelper {
    public static final String TRINO_INSTANCE                   = "trino_instance";
    public static final String TRINO_CATALOG                    = "trino_catalog";
    public static final String TRINO_SCHEMA                     = "trino_schema";
    public static final String TRINO_TABLE                      = "trino_table";
    public static final String TRINO_COLUMN                     = "trino_column";
    public static final String TRINO_INSTANCE_CATALOG_ATTRIBUTE = "catalogs";
    public static final String TRINO_CATALOG_SCHEMA_ATTRIBUTE   = "schemas";
    public static final String TRINO_SCHEMA_TABLE_ATTRIBUTE     = "tables";
    public static final String QUALIFIED_NAME_ATTRIBUTE         = "qualifiedName";
    public static final String NAME_ATTRIBUTE                   = "name";
    public static final int    DEFAULT_PAGE_LIMIT               = 10000;
    private static final Logger LOG = LoggerFactory.getLogger(AtlasClientHelper.class);
    private static final String DEFAULT_ATLAS_URL                        = "http://localhost:21000/";
    private static final String APPLICATION_PROPERTY_ATLAS_ENDPOINT      = "atlas.rest.address";
    private static final String TRINO_CATALOG_CONNECTOR_TYPE_ATTRIBUTE   = "connectorType";
    private static final String TRINO_CATALOG_INSTANCE_ATTRIBUTE         = "instance";
    private static final String TRINO_CATALOG_INSTANCE_RELATIONSHIP      = "trino_instance_catalog";
    private static final String TRINO_SCHEMA_CATALOG_ATTRIBUTE           = "catalog";
    private static final String TRINO_SCHEMA_CATALOG_RELATIONSHIP        = "trino_schema_catalog";
    private static final String TRINO_COLUMN_DATA_TYPE_ATTRIBUTE         = "data_type";
    private static final String TRINO_COLUMN_ORIDINAL_POSITION_ATTRIBUTE = "ordinal_position";
    private static final String TRINO_COLUMN_COLUMN_DEFAULT_ATTRIBUTE    = "column_default";
    private static final String TRINO_COLUMN_IS_NULLABLE_ATTRIBUTE       = "is_nullable";
    private static final String TRINO_COLUMN_TABLE_ATTRIBUTE             = "table";
    private static final String TRINO_TABLE_TYPE                         = "table_type";
    private static final String TRINO_TABLE_COLUMN_RELATIONSHIP          = "trino_table_columns";
    private static final String TRINO_TABLE_SCHEMA_RELATIONSHIP          = "trino_table_schema";
    private static final String TRINO_TABLE_SCHEMA_ATTRIBUTE             = "trinoschema";
    private static final String TRINO_TABLE_COLUMN_ATTRIBUTE             = "columns";

    private static AtlasClientV2 atlasClientV2;

    public AtlasClientHelper(Configuration atlasConf) throws IOException {
        atlasClientV2 = getAtlasClientV2Instance(atlasConf);
    }

    public List<AtlasEntityHeader> getAllCatalogsInInstance(String instanceGuid) throws AtlasServiceException {
        List<AtlasEntityHeader> entities = getAllRelationshipEntities(instanceGuid, TRINO_INSTANCE_CATALOG_ATTRIBUTE);

        if (CollectionUtils.isNotEmpty(entities)) {
            LOG.debug("Retrieved {} catalogs in Trino instance {}", entities.size(), instanceGuid);
        } else {
            LOG.debug("No catalog found in Trino instance {}", instanceGuid);
        }

        return entities;
    }

    public List<AtlasEntityHeader> getAllSchemasInCatalog(String catalogGuid) throws AtlasServiceException {
        List<AtlasEntityHeader> entities = getAllRelationshipEntities(catalogGuid, TRINO_CATALOG_SCHEMA_ATTRIBUTE);

        if (CollectionUtils.isNotEmpty(entities)) {
            LOG.debug("Retrieved {} schemas in Trino catalog {}", entities.size(), catalogGuid);
        } else {
            LOG.debug("No schema found in Trino catalog {}", catalogGuid);
        }

        return entities;
    }

    public List<AtlasEntityHeader> getAllTablesInSchema(String schemaGuid) throws AtlasServiceException {
        List<AtlasEntityHeader> entities = getAllRelationshipEntities(schemaGuid, TRINO_SCHEMA_TABLE_ATTRIBUTE);

        if (CollectionUtils.isNotEmpty(entities)) {
            LOG.debug("Retrieved {} tables in Trino schema {}", entities.size(), schemaGuid);
        } else {
            LOG.debug("No table found in Trino schema {}", schemaGuid);
        }

        return entities;
    }

    public AtlasEntityHeader getTrinoInstance(String namespace) {
        try {
            return atlasClientV2.getEntityHeaderByAttribute(TRINO_INSTANCE, Collections.singletonMap(QUALIFIED_NAME_ATTRIBUTE, namespace));
        } catch (AtlasServiceException e) {
            return null;
        }
    }

    public AtlasEntityWithExtInfo createOrUpdateInstanceEntity(String trinoNamespace) throws AtlasServiceException {
        String                 qualifiedName = trinoNamespace;
        AtlasEntityWithExtInfo ret           = findEntity(TRINO_INSTANCE, qualifiedName, true, true);

        if (ret == null) {
            AtlasEntity entity = new AtlasEntity(TRINO_INSTANCE);

            entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, qualifiedName);
            entity.setAttribute(NAME_ATTRIBUTE, trinoNamespace);

            ret = createEntity(new AtlasEntityWithExtInfo(entity));
        }

        return ret;
    }

    public AtlasEntityWithExtInfo createOrUpdateCatalogEntity(Catalog catalog) throws AtlasServiceException {
        String catalogName    = catalog.getName();
        String trinoNamespace = catalog.getInstanceName();
        String qualifiedName  = String.format("%s@%s", catalogName, trinoNamespace);

        AtlasEntityWithExtInfo ret = findEntity(TRINO_CATALOG, qualifiedName, true, true);

        if (ret == null) {
            AtlasEntity entity = new AtlasEntity(TRINO_CATALOG);

            entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, qualifiedName);
            entity.setAttribute(NAME_ATTRIBUTE, catalogName);
            entity.setAttribute(TRINO_CATALOG_CONNECTOR_TYPE_ATTRIBUTE, catalog.getType());
            entity.setRelationshipAttribute(TRINO_CATALOG_INSTANCE_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(catalog.getTrinoInstanceEntity().getEntity(), TRINO_CATALOG_INSTANCE_RELATIONSHIP));

            if (catalog.getConnector() != null) {
                catalog.getConnector().connectTrinoCatalog(this, catalog.getHookInstanceName(), catalogName, entity, ret);
            }

            ret = createEntity(new AtlasEntityWithExtInfo(entity));
        } else {
            AtlasEntity entity = ret.getEntity();

            entity.setRelationshipAttribute(TRINO_CATALOG_INSTANCE_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(catalog.getTrinoInstanceEntity().getEntity(), TRINO_CATALOG_INSTANCE_RELATIONSHIP));

            if (catalog.getConnector() != null) {
                catalog.getConnector().connectTrinoCatalog(this, catalog.getHookInstanceName(), catalogName, entity, ret);
            }

            ret.setEntity(entity);

            updateEntity(ret);
        }

        return ret;
    }

    public AtlasEntityWithExtInfo createOrUpdateSchemaEntity(Catalog catalog, AtlasEntity catalogEntity, String schema) throws AtlasServiceException {
        String                 qualifiedName = String.format("%s.%s@%s", catalog.getName(), schema, catalog.getInstanceName());
        AtlasEntityWithExtInfo ret           = findEntity(TRINO_SCHEMA, qualifiedName, true, true);

        if (ret == null) {
            AtlasEntity entity = new AtlasEntity(TRINO_SCHEMA);

            entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, qualifiedName);
            entity.setAttribute(NAME_ATTRIBUTE, schema);
            entity.setRelationshipAttribute(TRINO_SCHEMA_CATALOG_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(catalogEntity, TRINO_SCHEMA_CATALOG_RELATIONSHIP));

            if (catalog.getConnector() != null) {
                catalog.getConnector().connectTrinoSchema(this, catalog.getHookInstanceName(), catalog.getName(), schema, entity, ret);
            }

            ret = createEntity(new AtlasEntityWithExtInfo(entity));
        } else {
            AtlasEntity entity = ret.getEntity();

            entity.setRelationshipAttribute(TRINO_SCHEMA_CATALOG_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(catalogEntity, TRINO_SCHEMA_CATALOG_RELATIONSHIP));

            if (catalog.getConnector() != null) {
                catalog.getConnector().connectTrinoSchema(this, catalog.getHookInstanceName(), catalog.getName(), schema, entity, ret);
            }

            ret.setEntity(entity);

            updateEntity(ret);
        }

        return ret;
    }

    public AtlasEntityWithExtInfo createOrUpdateTableEntity(Catalog catalog, String schema, String table, Map<String, Object> tableMetadata, Map<String, Map<String, Object>> trinoColumns, AtlasEntity schemaEntity) throws AtlasServiceException {
        String qualifiedName = String.format("%s.%s.%s@%s", catalog.getName(), schema, table, catalog.getInstanceName());

        AtlasEntityWithExtInfo ret;
        AtlasEntityWithExtInfo tableEntityExt = findEntity(TRINO_TABLE, qualifiedName, true, true);

        if (tableEntityExt == null) {
            tableEntityExt = toTableEntity(catalog, schema, table, tableMetadata, trinoColumns, schemaEntity, tableEntityExt);
            ret            = createEntity(tableEntityExt);
        } else {
            ret = toTableEntity(catalog, schema, table, tableMetadata, trinoColumns, schemaEntity, tableEntityExt);

            updateEntity(ret);
        }

        return ret;
    }

    public void deleteByGuid(Set<String> guidTodelete) throws AtlasServiceException {
        if (CollectionUtils.isNotEmpty(guidTodelete)) {
            for (String guid : guidTodelete) {
                EntityMutationResponse response = atlasClientV2.deleteEntityByGuid(guid);

                if (response == null || response.getDeletedEntities().isEmpty()) {
                    LOG.debug("Entity with guid : {} is not deleted", guid);
                } else {
                    LOG.debug("Entity with guid : {} is deleted", guid);
                }
            }
        }
    }

    public AtlasEntityWithExtInfo findEntity(final String typeName, final String qualifiedName, boolean minExtInfo, boolean ignoreRelationship) throws AtlasServiceException {
        try {
            return atlasClientV2.getEntityByAttribute(typeName, Collections.singletonMap(ATTRIBUTE_QUALIFIED_NAME, qualifiedName), minExtInfo, ignoreRelationship);
        } catch (AtlasServiceException e) {
            if (e.getStatus() == ClientResponse.Status.NOT_FOUND) {
                return null;
            }

            throw e;
        }
    }

    public void close() {
        if (atlasClientV2 != null) {
            atlasClientV2.close();

            atlasClientV2 = null;
        }
    }

    private synchronized AtlasClientV2 getAtlasClientV2Instance(Configuration atlasConf) throws IOException {
        if (atlasClientV2 == null) {
            String[] atlasEndpoint = new String[] {DEFAULT_ATLAS_URL};

            if (atlasConf != null && ArrayUtils.isNotEmpty(atlasConf.getStringArray(APPLICATION_PROPERTY_ATLAS_ENDPOINT))) {
                atlasEndpoint = atlasConf.getStringArray(APPLICATION_PROPERTY_ATLAS_ENDPOINT);
            }

            if (!AuthenticationUtil.isKerberosAuthenticationEnabled()) {
                String[] basicAuthUsernamePassword = AuthenticationUtil.getBasicAuthenticationInput();

                atlasClientV2 = new AtlasClientV2(atlasEndpoint, basicAuthUsernamePassword);
            } else {
                UserGroupInformation ugi = UserGroupInformation.getCurrentUser();

                atlasClientV2 = new AtlasClientV2(ugi, ugi.getShortUserName(), atlasEndpoint);
            }
        }

        return atlasClientV2;
    }

    private List<AtlasEntityHeader> getAllRelationshipEntities(String entityGuid, String relationshipAttributeName) throws AtlasServiceException {
        List<AtlasEntityHeader> entities = new ArrayList<>();

        if (entityGuid != null) {
            final int pageSize = DEFAULT_PAGE_LIMIT;

            for (int i = 0; ; i++) {
                int offset = pageSize * i;

                LOG.debug("Retrieving: offset={}, pageSize={}", offset, pageSize);

                AtlasSearchResult       searchResult  = atlasClientV2.relationshipSearch(entityGuid, relationshipAttributeName, null, null, true, pageSize, offset);
                List<AtlasEntityHeader> entityHeaders = searchResult == null ? null : searchResult.getEntities();
                int                     count         = entityHeaders == null ? 0 : entityHeaders.size();

                if (count > 0) {
                    entities.addAll(entityHeaders);
                }

                if (count < pageSize) { // last page
                    break;
                }
            }
        }

        return entities;
    }

    private AtlasEntityWithExtInfo toTableEntity(Catalog catalog, String schema, String table, Map<String, Object> tableMetadata, Map<String, Map<String, Object>> trinoColumns, AtlasEntity schemaEntity, AtlasEntityWithExtInfo tableEntityExt) {
        if (tableEntityExt == null) {
            tableEntityExt = new AtlasEntityWithExtInfo(new AtlasEntity(TRINO_TABLE));
        }

        AtlasEntity       tableEntity    = tableEntityExt.getEntity();
        List<AtlasEntity> columnEntities = new ArrayList<>();
        String            qualifiedName  = String.format("%s.%s.%s@%s", catalog.getName(), schema, table, catalog.getInstanceName());

        tableEntity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, qualifiedName);
        tableEntity.setAttribute(NAME_ATTRIBUTE, table);
        tableEntity.setAttribute(TRINO_TABLE_TYPE, tableMetadata.get(TRINO_TABLE_TYPE).toString());

        for (Map.Entry<String, Map<String, Object>> columnEntry : trinoColumns.entrySet()) {
            AtlasEntity entity           = new AtlasEntity(TRINO_COLUMN);
            String      columnName       = columnEntry.getKey();
            String      colQualifiedName = String.format("%s.%s.%s.%s@%s", catalog.getName(), schema, table, columnName, catalog.getInstanceName());

            entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, colQualifiedName);
            entity.setAttribute(NAME_ATTRIBUTE, columnName);

            if (MapUtils.isNotEmpty(columnEntry.getValue())) {
                Map<String, Object> columnAttr = columnEntry.getValue();

                entity.setAttribute(TRINO_COLUMN_DATA_TYPE_ATTRIBUTE, columnAttr.get(TRINO_COLUMN_DATA_TYPE_ATTRIBUTE));
                entity.setAttribute(TRINO_COLUMN_ORIDINAL_POSITION_ATTRIBUTE, columnAttr.get(TRINO_COLUMN_ORIDINAL_POSITION_ATTRIBUTE));
                entity.setAttribute(TRINO_COLUMN_COLUMN_DEFAULT_ATTRIBUTE, columnAttr.get(TRINO_COLUMN_COLUMN_DEFAULT_ATTRIBUTE));
                entity.setAttribute(TRINO_COLUMN_IS_NULLABLE_ATTRIBUTE, columnAttr.get(TRINO_COLUMN_IS_NULLABLE_ATTRIBUTE));
            }

            entity.setRelationshipAttribute(TRINO_COLUMN_TABLE_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(tableEntity, TRINO_TABLE_COLUMN_RELATIONSHIP));

            columnEntities.add(entity);
        }

        tableEntity.setRelationshipAttribute(TRINO_TABLE_SCHEMA_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectId(schemaEntity, TRINO_TABLE_SCHEMA_RELATIONSHIP));
        tableEntity.setRelationshipAttribute(TRINO_TABLE_COLUMN_ATTRIBUTE, AtlasTypeUtil.getAtlasRelatedObjectIds(columnEntities, TRINO_TABLE_COLUMN_RELATIONSHIP));

        if (catalog.getConnector() != null) {
            catalog.getConnector().connectTrinoTable(this, catalog.getHookInstanceName(), catalog.getName(), schema, table, tableEntity, columnEntities, tableEntityExt);
        }

        tableEntityExt.addReferredEntity(schemaEntity);

        for (AtlasEntity column : columnEntities) {
            tableEntityExt.addReferredEntity(column);
        }

        return tableEntityExt;
    }

    private AtlasEntityWithExtInfo createEntity(AtlasEntityWithExtInfo entity) throws AtlasServiceException {
        LOG.debug("creating {} entity: {}", entity.getEntity().getTypeName(), entity);

        AtlasEntityWithExtInfo  ret             = null;
        EntityMutationResponse  response        = atlasClientV2.createEntity(entity);
        List<AtlasEntityHeader> createdEntities = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

        if (CollectionUtils.isNotEmpty(createdEntities)) {
            for (AtlasEntityHeader createdEntity : createdEntities) {
                if (ret == null) {
                    ret = atlasClientV2.getEntityByGuid(createdEntity.getGuid());

                    LOG.debug("Created {} entity: name={}, guid={}", ret.getEntity().getTypeName(), ret.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), ret.getEntity().getGuid());
                } else if (ret.getEntity(createdEntity.getGuid()) == null) {
                    AtlasEntityWithExtInfo newEntity = atlasClientV2.getEntityByGuid(createdEntity.getGuid());

                    ret.addReferredEntity(newEntity.getEntity());

                    if (MapUtils.isNotEmpty(newEntity.getReferredEntities())) {
                        for (Map.Entry<String, AtlasEntity> entry : newEntity.getReferredEntities().entrySet()) {
                            ret.addReferredEntity(entry.getKey(), entry.getValue());
                        }
                    }

                    LOG.debug("Created {} entity: name={}, guid={}", newEntity.getEntity().getTypeName(), newEntity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), newEntity.getEntity().getGuid());
                }
            }
        }

        clearRelationshipAttributes(ret);

        return ret;
    }

    private void updateEntity(AtlasEntityWithExtInfo entity) throws AtlasServiceException {
        LOG.debug("updating {} entity: {}", entity.getEntity().getTypeName(), entity);

        atlasClientV2.updateEntity(entity);

        LOG.debug("Updated {} entity: name={}, guid={}", entity.getEntity().getTypeName(), entity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), entity.getEntity().getGuid());
    }

    private void clearRelationshipAttributes(AtlasEntityWithExtInfo entity) {
        if (entity != null) {
            clearRelationshipAttributes(entity.getEntity());

            if (entity.getReferredEntities() != null) {
                clearRelationshipAttributes(entity.getReferredEntities().values());
            }
        }
    }

    private void clearRelationshipAttributes(Collection<AtlasEntity> entities) {
        if (entities != null) {
            for (AtlasEntity entity : entities) {
                clearRelationshipAttributes(entity);
            }
        }
    }

    private void clearRelationshipAttributes(AtlasEntity entity) {
        if (entity != null && entity.getRelationshipAttributes() != null) {
            entity.getRelationshipAttributes().clear();
        }
    }
}
