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

import com.datastax.driver.core.BatchStatement;
import com.datastax.driver.core.BoundStatement;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.Session;
import com.google.common.annotations.VisibleForTesting;
import org.apache.atlas.AtlasException;
import org.apache.atlas.EntityAuditEvent;
import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.audit.EntityAuditEventV2;
import org.apache.atlas.repository.Constants;
import org.apache.commons.lang3.NotImplementedException;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.inject.Singleton;

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

/**
 * This class provides cassandra support as the backend for audit storage support.
 */
@Singleton
@Component
@ConditionalOnAtlasProperty(property = "atlas.EntityAuditRepository.impl")
public class CassandraBasedAuditRepository extends AbstractStorageBasedAuditRepository {
    private static final Logger LOG = LoggerFactory.getLogger(CassandraBasedAuditRepository.class);

    // The environment variable that tells us we are running in embedded mode
    public static final String MANAGE_EMBEDDED_CASSANDRA  = "MANAGE_EMBEDDED_CASSANDRA";

    // Application properties
    public static final String CASSANDRA_HOSTNAME_PROPERTY           = "atlas.graph.storage.hostname";
    public static final String CASSANDRA_CLUSTERNAME_PROPERTY        = "atlas.graph.storage.clustername";
    public static final String CASSANDRA_PORT_PROPERTY               = "atlas.graph.storage.port";
    public static final String CASSANDRA_REPLICATION_FACTOR_PROPERTY = "atlas.EntityAuditRepository.replicationFactor";
    public static final String CASSANDRA_AUDIT_KEYSPACE_PROPERTY     = "atlas.EntityAuditRepository.keyspace";
    public static final String CASSANDRA_USERNAME_PROPERTY           = "atlas.graph.storage.username";
    public static final String CASSANDRA_PASSWORD_PROPERTY           = "atlas.graph.storage.password";

    // Default keyspace to store the audit entries
    private static final String DEFAULT_KEYSPACE           = "atlas_audit";

    // When running in embedded cassandra mode, this is the default cluster name used
    private static final String DEFAULT_CLUSTER_NAME       = "JanusGraph";

    // Default cassandra port
    private static final int    DEFAULT_PORT               = 9042;
    private static final int    DEFAULT_REPLICATION_FACTOR = 3;
    private static final String AUDIT_TABLE_SCHEMA =
            "CREATE TABLE audit(entityid text, "
                    + "created bigint, "
                    + "action text, "
                    + "user text, "
                    + "detail text, "
                    + "entity text, "
                    + "PRIMARY KEY (entityid, created)"
                    + ") WITH CLUSTERING ORDER BY (created DESC);";

    private static final String ENTITYID = "entityid";
    private static final String CREATED  = "created";
    private static final String ACTION   = "action";
    private static final String USER     = "user";
    private static final String DETAIL   = "detail";
    private static final String ENTITY   = "entity";

    private static final String INSERT_STATEMENT_TEMPLATE      = "INSERT INTO audit (entityid,created,action,user,detail,entity) VALUES (?,?,?,?,?,?)";
    private static final String SELECT_STATEMENT_TEMPLATE      = "select * from audit where entityid=? order by created desc limit 10;";
    private static final String SELECT_DATE_STATEMENT_TEMPLATE = "select * from audit where entityid=? and created<=? order by created desc limit 10;";

    private String  keyspace;
    private int     replicationFactor;
    private Session cassSession;
    private String  clusterName;
    private int     port;
    private String  username;
    private String  password;

    private final Map<String, List<String>> auditExcludedAttributesCache = new HashMap<>();

    private PreparedStatement insertStatement;
    private PreparedStatement selectStatement;
    private PreparedStatement selectDateStatement;

    @Override
    public void putEventsV1(List<EntityAuditEvent> events) {
        BoundStatement stmt  = new BoundStatement(insertStatement);
        BatchStatement batch = new BatchStatement();

        events.forEach(event -> batch.add(stmt.bind(event.getEntityId(), event.getTimestamp(),
                event.getAction().toString(), event.getUser(), event.getDetails(),
                (PERSIST_ENTITY_DEFINITION ? event.getEntityDefinitionString() : null))));

        cassSession.execute(batch);
    }

    @Override
    public List<EntityAuditEvent> listEventsV1(String entityId, String startKey, short maxResults) {
        LOG.debug("Listing events for entity id {}, starting timestamp {}, #records {}", entityId, startKey, maxResults);

        ResultSet              rs            = cassSession.execute(getSelectStatement(entityId, startKey));
        List<EntityAuditEvent> entityResults = new ArrayList<>();

        for (Row row : rs) {
            String rowEntityId = row.getString(ENTITYID);

            if (!entityId.equals(rowEntityId)) {
                continue;
            }

            EntityAuditEvent event = new EntityAuditEvent();

            event.setEntityId(rowEntityId);
            event.setAction(EntityAuditEvent.EntityAuditAction.fromString(row.getString(ACTION)));
            event.setDetails(row.getString(DETAIL));
            event.setUser(row.getString(USER));
            event.setTimestamp(row.getLong(CREATED));
            event.setEventKey(rowEntityId + ":" + event.getTimestamp());

            if (PERSIST_ENTITY_DEFINITION) {
                event.setEntityDefinition(row.getString(ENTITY));
            }

            entityResults.add(event);
        }

        return entityResults;
    }

    @Override
    public void putEventsV2(List<EntityAuditEventV2> events) {
        BoundStatement stmt  = new BoundStatement(insertStatement);
        BatchStatement batch = new BatchStatement();

        events.forEach(event -> batch.add(stmt.bind(event.getEntityId(), event.getTimestamp(),
                event.getAction().toString(), event.getUser(), event.getDetails(),
                (PERSIST_ENTITY_DEFINITION ? event.getEntityDefinitionString() : null))));

        cassSession.execute(batch);
    }

    @Override
    public List<EntityAuditEventV2> listEventsV2(String entityId, EntityAuditEventV2.EntityAuditActionV2 auditAction, String startKey, short maxResults) {
        LOG.debug("Listing events for entity id {}, starting timestamp {}, #records {}", entityId, startKey, maxResults);

        ResultSet                rs            = cassSession.execute(getSelectStatement(entityId, startKey));
        List<EntityAuditEventV2> entityResults = new ArrayList<>();

        for (Row row : rs) {
            String rowEntityId = row.getString(ENTITYID);

            if (!entityId.equals(rowEntityId)) {
                continue;
            }

            EntityAuditEventV2 event = new EntityAuditEventV2();

            event.setEntityId(rowEntityId);
            event.setAction(EntityAuditEventV2.EntityAuditActionV2.fromString(row.getString(ACTION)));
            event.setDetails(row.getString(DETAIL));
            event.setUser(row.getString(USER));
            event.setTimestamp(row.getLong(CREATED));
            event.setEventKey(rowEntityId + ":" + event.getTimestamp());

            if (PERSIST_ENTITY_DEFINITION) {
                event.setEntityDefinition(row.getString(ENTITY));
            }

            entityResults.add(event);
        }

        return entityResults;
    }

    @Override
    public List<EntityAuditEventV2> listEventsV2(String entityId, EntityAuditEventV2.EntityAuditActionV2 auditAction, String sortByColumn, boolean sortOrderDesc, int offset, short limit) {
        return listEventsV2(entityId, auditAction, null, limit);
    }

    @Override
    public List<EntityAuditEventV2> deleteEventsV2(String entityId, Set<EntityAuditEventV2.EntityAuditActionV2> entityAuditActions, short auditCount, int ttlInDays, boolean createEventsAgeoutAllowed, Constants.AtlasAuditAgingType auditAgingType)
            throws AtlasBaseException, AtlasException {
        return Collections.emptyList();
    }

    @Override
    public Set<String> getEntitiesWithTagChanges(long fromTimestamp, long toTimestamp) {
        throw new NotImplementedException();
    }

    @Override
    public void start() throws AtlasException {
        initApplicationProperties();
        initializeSettings();
        startInternal();
    }

    @Override
    public void stop() throws AtlasException {
        cassSession.close();
    }

    void initializeSettings() {
        keyspace          = applicationProperties.getString(CASSANDRA_AUDIT_KEYSPACE_PROPERTY, DEFAULT_KEYSPACE);
        replicationFactor = applicationProperties.getInt(CASSANDRA_REPLICATION_FACTOR_PROPERTY, DEFAULT_REPLICATION_FACTOR);
        clusterName       = applicationProperties.getString(CASSANDRA_CLUSTERNAME_PROPERTY, DEFAULT_CLUSTER_NAME);
        port              = applicationProperties.getInt(CASSANDRA_PORT_PROPERTY, DEFAULT_PORT);
        username          = applicationProperties.getString(CASSANDRA_USERNAME_PROPERTY, "");
        password          = applicationProperties.getString(CASSANDRA_PASSWORD_PROPERTY, "");
    }

    @VisibleForTesting
    void startInternal() throws AtlasException {
        createSession();
    }

    void createSession() throws AtlasException {
        Cluster.Builder cassandraClusterBuilder = Cluster.builder();

        String  hostname = applicationProperties.getString(CASSANDRA_HOSTNAME_PROPERTY, "localhost");
        Cluster cluster  = cassandraClusterBuilder.addContactPoint(hostname).withClusterName(clusterName).withPort(port).withCredentials(username.trim(), password.trim()).build();

        try {
            cassSession = cluster.connect();

            if (cluster.getMetadata().getKeyspace(keyspace) == null) {
                String query = "CREATE KEYSPACE " + keyspace + " WITH replication = {'class':'SimpleStrategy', 'replication_factor':" + replicationFactor + "}; ";

                cassSession.execute(query);
                cassSession.close();

                cassSession = cluster.connect(keyspace);

                // create the audit table
                cassSession.execute(AUDIT_TABLE_SCHEMA);
            } else {
                cassSession.close();
                cassSession = cluster.connect(keyspace);
            }

            insertStatement     = cassSession.prepare(INSERT_STATEMENT_TEMPLATE.replace("KEYSPACE", keyspace));
            selectStatement     = cassSession.prepare(SELECT_STATEMENT_TEMPLATE.replace("KEYSPACE", keyspace));
            selectDateStatement = cassSession.prepare(SELECT_DATE_STATEMENT_TEMPLATE.replace("KEYSPACE", keyspace));
        } catch (Exception e) {
            throw new AtlasException(e);
        }
    }

    private BoundStatement getSelectStatement(String entityId, String startKey) {
        BoundStatement stmt;

        if (StringUtils.isEmpty(startKey)) {
            stmt = new BoundStatement(selectStatement).bind(entityId);
        } else {
            stmt = new BoundStatement(selectDateStatement).bind(entityId, Long.valueOf(startKey.split(FIELD_SEPARATOR)[1]));
        }

        return stmt;
    }
}
