/*
 * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
 *
 * WSO2 LLC. 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
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.wso2.carbon.is.migration.service.v710.dao;

import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.wso2.carbon.context.CarbonContext;
import org.wso2.carbon.database.utils.jdbc.NamedPreparedStatement;
import org.wso2.carbon.identity.core.migrate.MigrationClientException;
import org.wso2.carbon.identity.core.util.IdentityDatabaseUtil;
import org.wso2.carbon.identity.entitlement.EntitlementUtil;
import org.wso2.carbon.identity.entitlement.common.EntitlementConstants;
import org.wso2.carbon.identity.entitlement.dto.AttributeDTO;
import org.wso2.carbon.identity.entitlement.dto.PolicyDTO;
import org.wso2.carbon.identity.entitlement.dto.PublisherDataHolder;
import org.wso2.carbon.identity.entitlement.dto.PublisherPropertyDTO;
import org.wso2.carbon.identity.entitlement.dto.StatusHolder;
import org.wso2.carbon.is.migration.service.v710.constants.XACMLRegistryConstants;
import org.wso2.carbon.is.migration.service.v710.model.PolicyMigrationDTO;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.Map;
import java.util.TimeZone;

import static org.wso2.carbon.is.migration.service.v710.constants.XACMLRegistryConstants.GLOBAL_POLICY_COMBINING_ALGORITHM;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_ATTRIBUTES_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_COMBINING_ALGORITHM_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_EDITOR_DATA_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_REFS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_SET_REFS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_POLICY_STATUS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_SUBSCRIBER_PROPERTIES_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_SUBSCRIBER_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.CREATE_SUBSCRIBER_STATUS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.GET_POLICY_COMBINING_ALGORITHM_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.LIST_POLICY_IDS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.LIST_POLICY_IDS_WITH_STATUS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.LIST_SUBSCRIBER_IDS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.LIST_SUBSCRIBER_IDS_WITH_STATUS_SQL;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.ATTRIBUTE_ID;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.ATTRIBUTE_VALUE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.CATEGORY;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.CONFIG_KEY;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.CONFIG_VALUE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.DATA_TYPE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.DISPLAY_NAME;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.DISPLAY_ORDER;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.EDITOR_DATA;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.EDITOR_DATA_ORDER;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.ENTITLEMENT_MODULE_NAME;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_ACTIVE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_IN_PAP;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_IN_PDP;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_REQUIRED;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_SECRET;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.IS_SUCCESS;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.KEY;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.LAST_MODIFIED_TIME;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.LAST_MODIFIED_USER;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.LOGGED_AT;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.MESSAGE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.MODULE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.POLICY;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.POLICY_EDITOR;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.POLICY_ID;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.POLICY_ORDER;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.POLICY_TYPE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.PROPERTY_ID;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.PROPERTY_VALUE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.REFERENCE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.SET_REFERENCE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.STATUS_TYPE;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.SUBSCRIBER_ID;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.TARGET;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.TARGET_ACTION;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.TENANT_ID;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.USER;
import static org.wso2.carbon.is.migration.service.v710.constants.XACMLSQLConstants.XACMLTableColumns.VERSION;

import static java.time.ZoneOffset.UTC;

/**
 * This class handles the migrations related to XACML.
 */
public class XACMLDAO {

    private static final Logger LOG = LoggerFactory.getLogger(XACMLDAO.class);

    public void migrateXACMLConfigs(String xacmlConfig, boolean isDryRun, int tenantId) throws
            MigrationClientException {

        if (StringUtils.isBlank(xacmlConfig)) {
            LOG.info("RDM-XACML>> No XACML configs to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("RDM-XACML>> Migrating XACML configs in tenant: {}.", tenantId);

        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
        try {
            String algorithm = getPolicyCombiningAlgorithm(connection, tenantId);
            if (StringUtils.isNotBlank(algorithm)) {
                LOG.info("RDM-XACML>> Global policy combining algorithm {} already exists in tenant: {}.",
                        algorithm, tenantId);
            } else {
                LOG.info("RDM-XACML>> Global policy combining algorithm {} selected for migration in tenant: {}.",
                        xacmlConfig, tenantId);
                if (!isDryRun) {
                    storePolicyCombiningAlgorithm(connection, xacmlConfig, tenantId);
                } else {
                    LOG.info("RDM-XACML>> Dry run enabled. Skipping the global policy combining algorithm " +
                            "migration in tenant: {}.", tenantId);
                }
            }
            IdentityDatabaseUtil.commitTransaction(connection);
        } catch (SQLException e) {
            IdentityDatabaseUtil.rollbackTransaction(connection);
            throw new MigrationClientException(String.format("Error while migrating the global policy combining " +
                    "algorithm %s in tenant: %s.", xacmlConfig, tenantId), e);
        } finally {
            IdentityDatabaseUtil.closeConnection(connection);
        }
    }

    public void migrateXACMLSubscribers(List<PublisherDataHolder> xacmlSubscribers, boolean isDryRun, int tenantId)
            throws MigrationClientException {

        if (xacmlSubscribers == null || xacmlSubscribers.isEmpty()) {
            LOG.info("RDM-XACML>> No XACML subscribers to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("RDM-XACML>> {} XACML subscribers picked for migration from Registry in tenant: {}.",
                xacmlSubscribers.size(), tenantId);

        // Perform the migration in a transaction
        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
        try {
            // Get all existing subscribers in the DB
            List<String> existingSubscribers = getExistingSubscriberIds(connection, tenantId);
            if (!existingSubscribers.isEmpty()) {
                LOG.info("RDM-XACML>> {} existing XACML subscribers found in the DB in tenant: {}.",
                        existingSubscribers.size(), tenantId);
                // Filter out the subscribers that are already in the DB
                xacmlSubscribers.removeIf(subscriber -> {
                    PublisherPropertyDTO[] propertyDTOs = subscriber.getPropertyDTOs();
                    for (PublisherPropertyDTO propertyDTO : propertyDTOs) {
                        if (XACMLRegistryConstants.SUBSCRIBER_ID.equals(propertyDTO.getId())) {
                            return existingSubscribers.contains(propertyDTO.getValue());
                        }
                    }
                    return false;
                });
            }
            LOG.info("RDM-XACML>> {} subscribers selected for migration in tenant: {}.", xacmlSubscribers.size(),
                    tenantId);
            // Add the filtered subscribers to the DB
            if (!isDryRun) {
                storeSubscribers(connection, xacmlSubscribers, tenantId);
            } else {
                LOG.info("RDM-XACML>> Dry run enabled. Skipping the XACML subscribers migration in tenant: {}.",
                        tenantId);
            }
            IdentityDatabaseUtil.commitTransaction(connection);
        } catch (SQLException e) {
            IdentityDatabaseUtil.rollbackTransaction(connection);
            throw new MigrationClientException(String.format("Error while migrating subscribers in tenant: %s.",
                    tenantId), e);
        } finally {
            IdentityDatabaseUtil.closeConnection(connection);
        }
    }

    public void migrateXACMLPolicyStatuses(Map<String, List<StatusHolder>> policyStatusHolders, boolean isDryRun,
                                           int tenantId) throws MigrationClientException {

        if (policyStatusHolders == null || policyStatusHolders.isEmpty()) {
            LOG.info("RDM-XACML>> No XACML statuses to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("RDM-XACML>> {} XACML statuses picked for migration from Registry in tenant: {}.",
                policyStatusHolders.size(), tenantId);

        // Perform the migration in a transaction
        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
        try {
            // Get all existing policies for which a status trail exists in the DB
            List<String> existingPolicyIdsWithStatuses = getPolicyIdsWithStatusTrail(connection, tenantId);
            if (!existingPolicyIdsWithStatuses.isEmpty()) {
                LOG.info("RDM-XACML>> {} existing XACML status trails found in the DB for policies in tenant: {}.",
                        existingPolicyIdsWithStatuses.size(), tenantId);
                // Filter out the status trails of policies for which a trail already exists in DB
                policyStatusHolders.keySet().removeIf(existingPolicyIdsWithStatuses::contains);
            }

            // Add the filtered statuses to the DB
            if (!isDryRun) {
                storeStatuses(connection, policyStatusHolders, XACMLRegistryConstants.POLICY, tenantId);
            } else {
                LOG.info("RDM-XACML>> Dry run enabled. Skipping the XACML subscribers migration in tenant: {}.",
                        tenantId);
            }
            IdentityDatabaseUtil.commitTransaction(connection);
        } catch (SQLException e) {
            IdentityDatabaseUtil.rollbackTransaction(connection);
            throw new MigrationClientException(String.format("Error while migrating subscribers in tenant: %s.",
                    tenantId), e);
        } finally {
            IdentityDatabaseUtil.closeConnection(connection);
        }
    }

    public void migrateXACMLSubscriberStatuses(Map<String, List<StatusHolder>> subscriberStatusHolders,
                                               boolean isDryRun, int tenantId) throws MigrationClientException {

        if (subscriberStatusHolders == null || subscriberStatusHolders.isEmpty()) {
            LOG.info("RDM-XACML>> No XACML statuses to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("RDM-XACML>> {} XACML statuses picked for migration from Registry in tenant: {}.",
                subscriberStatusHolders.size(), tenantId);

        // Perform the migration in a transaction
        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
        try {
            // Get all existing subscribers for which a status trail exists in the DB
            List<String> existingSubscriberIdsWithStatuses = getSubscriberIdsWithStatusTrail(connection, tenantId);
            if (!existingSubscriberIdsWithStatuses.isEmpty()) {
                LOG.info("RDM-XACML>> {} existing XACML status trails found in the DB for subscribers in tenant: {}.",
                        existingSubscriberIdsWithStatuses.size(), tenantId);
                // Filter out the status trails of subscribers for which a trail already exists in DB
                subscriberStatusHolders.keySet().removeIf(existingSubscriberIdsWithStatuses::contains);
            }

            // Add the filtered statuses to the DB
            if (!isDryRun) {
                storeStatuses(connection, subscriberStatusHolders, XACMLRegistryConstants.SUBSCRIBER, tenantId);
            } else {
                LOG.info("RDM-XACML>> Dry run enabled. Skipping the XACML subscribers migration in tenant: {}.",
                        tenantId);
            }
            IdentityDatabaseUtil.commitTransaction(connection);
        } catch (SQLException e) {
            IdentityDatabaseUtil.rollbackTransaction(connection);
            throw new MigrationClientException(String.format("Error while migrating subscribers in tenant: %s.",
                    tenantId), e);
        } finally {
            IdentityDatabaseUtil.closeConnection(connection);
        }
    }

    public void migrateXACMLPolicies(List<PolicyMigrationDTO> xacmlPolicies, boolean isDryRun, int tenantId)
            throws MigrationClientException {

        if (xacmlPolicies == null || xacmlPolicies.isEmpty()) {
            LOG.info("RDM-XACML>> No XACML policies to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("RDM-XACML>> {} XACML PAP policies picked for migration from Registry in tenant: {}.",
                xacmlPolicies.size(), tenantId);
        // Perform the migration in a transaction
        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
        try {
            // Get all existing policies in the DB
            List<String> existingPolicies = getExistingPolicyIds(connection, tenantId);
            if (!existingPolicies.isEmpty()) {
                LOG.info("RDM-XACML>> {} existing XACML policies found in the DB in tenant: {}.",
                        existingPolicies.size(), tenantId);
                // Filter out the policies that are already in the DB
                xacmlPolicies.removeIf(policy -> existingPolicies.contains(policy.getPolicyDTO().getPolicyId()));
            }
            LOG.info("RDM-XACML>> {} policies selected for migration in tenant: {}.", existingPolicies.size(),
                    tenantId);
            // Add the filtered policies to the DB
            if (!isDryRun) {
                storePolicies(connection, xacmlPolicies, tenantId);
            } else {
                LOG.info("RDM-XACML>> Dry run enabled. Skipping the XACML policies migration in tenant: {}.",
                        tenantId);
            }
            IdentityDatabaseUtil.commitTransaction(connection);
        } catch (SQLException e) {
            IdentityDatabaseUtil.rollbackTransaction(connection);
            throw new MigrationClientException(String.format("Error while migrating policies in tenant: %s.",
                    tenantId), e);
        } finally {
            IdentityDatabaseUtil.closeConnection(connection);
        }
    }

    private String getPolicyCombiningAlgorithm(Connection connection, int tenantId) throws SQLException {

        String algorithm = null;
        try (NamedPreparedStatement getPolicyCombiningAlgoPrepStmt = new NamedPreparedStatement(connection,
                GET_POLICY_COMBINING_ALGORITHM_SQL)) {
            getPolicyCombiningAlgoPrepStmt.setString(CONFIG_KEY, GLOBAL_POLICY_COMBINING_ALGORITHM);
            getPolicyCombiningAlgoPrepStmt.setInt(TENANT_ID, tenantId);
            try (ResultSet rs = getPolicyCombiningAlgoPrepStmt.executeQuery()) {
                if (rs.next()) {
                    algorithm = rs.getString(CONFIG_VALUE);
                }
            }
        }
        return algorithm;
    }

    private void storePolicyCombiningAlgorithm(Connection connection, String policyCombiningAlgorithm, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement setPolicyCombiningAlgoPrepStmt = new NamedPreparedStatement(connection,
                CREATE_POLICY_COMBINING_ALGORITHM_SQL)) {
            setPolicyCombiningAlgoPrepStmt.setString(CONFIG_KEY, GLOBAL_POLICY_COMBINING_ALGORITHM);
            setPolicyCombiningAlgoPrepStmt.setString(CONFIG_VALUE, policyCombiningAlgorithm);
            setPolicyCombiningAlgoPrepStmt.setInt(TENANT_ID, tenantId);
            setPolicyCombiningAlgoPrepStmt.executeUpdate();
        }
    }

    private void storeSubscribers(Connection connection, List<PublisherDataHolder> xacmlSubscribers, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement createSubscriberPrepStmt = new NamedPreparedStatement(connection,
                CREATE_SUBSCRIBER_SQL);
             NamedPreparedStatement createSubscriberPropertiesPrepStmt = new NamedPreparedStatement(connection,
                     CREATE_SUBSCRIBER_PROPERTIES_SQL)) {

            for (PublisherDataHolder holder : xacmlSubscribers) {
                String subscriberId = resolveSubscriberId(holder);
                if (StringUtils.isBlank(subscriberId)) {
                    continue;
                }
                createSubscriberPrepStmt.setString(SUBSCRIBER_ID, subscriberId);
                createSubscriberPrepStmt.setString(ENTITLEMENT_MODULE_NAME, holder.getModuleName());
                createSubscriberPrepStmt.setInt(TENANT_ID, tenantId);
                createSubscriberPrepStmt.addBatch();

                for (PublisherPropertyDTO dto : holder.getPropertyDTOs()) {
                    if (dto.getId() != null && StringUtils.isNotBlank(dto.getValue())) {

                        createSubscriberPropertiesPrepStmt.setString(PROPERTY_ID, dto.getId());
                        createSubscriberPropertiesPrepStmt.setString(DISPLAY_NAME, dto.getDisplayName());
                        createSubscriberPropertiesPrepStmt.setString(PROPERTY_VALUE, dto.getValue());
                        createSubscriberPropertiesPrepStmt.setBoolean(IS_REQUIRED, dto.isRequired());
                        createSubscriberPropertiesPrepStmt.setInt(DISPLAY_ORDER, dto.getDisplayOrder());
                        createSubscriberPropertiesPrepStmt.setBoolean(IS_SECRET, dto.isSecret());
                        createSubscriberPropertiesPrepStmt.setString(MODULE, dto.getModule());
                        createSubscriberPropertiesPrepStmt.setString(SUBSCRIBER_ID, subscriberId);
                        createSubscriberPropertiesPrepStmt.setInt(TENANT_ID, tenantId);

                        createSubscriberPropertiesPrepStmt.addBatch();
                    }
                }
            }
            createSubscriberPrepStmt.executeBatch();
            createSubscriberPropertiesPrepStmt.executeBatch();
        }
    }

    private static String resolveSubscriberId(PublisherDataHolder holder) {

        String subscriberId = null;
        if (holder == null || holder.getPropertyDTOs() == null) {
            return null;
        }

        for (PublisherPropertyDTO dto : holder.getPropertyDTOs()) {
            if (XACMLRegistryConstants.SUBSCRIBER_ID.equals(dto.getId())) {
                subscriberId = dto.getValue();
            }
        }
        return subscriberId;
    }

    private List<String> getExistingSubscriberIds(Connection connection, int tenantId) throws SQLException {

        List<String> subscriberIdList = new ArrayList<>();
        try (NamedPreparedStatement preparedStmt = new NamedPreparedStatement(connection, LIST_SUBSCRIBER_IDS_SQL)) {
            preparedStmt.setInt(TENANT_ID, tenantId);
            try (ResultSet subscriberIds = preparedStmt.executeQuery()) {
                while (subscriberIds.next()) {
                    subscriberIdList.add(subscriberIds.getString(SUBSCRIBER_ID));
                }
            }

        }
        return subscriberIdList;
    }

    private List<String> getPolicyIdsWithStatusTrail(Connection connection, int tenantId) throws SQLException {

        List<String> policyIdList = new ArrayList<>();
        try (NamedPreparedStatement preparedStmt = new NamedPreparedStatement(connection,
                LIST_POLICY_IDS_WITH_STATUS_SQL)) {
            preparedStmt.setInt(TENANT_ID, tenantId);
            try (ResultSet resultSet = preparedStmt.executeQuery()) {
                while (resultSet.next()) {
                    policyIdList.add(resultSet.getString(POLICY_ID));
                }
            }
        }
        return policyIdList;
    }

    private List<String> getSubscriberIdsWithStatusTrail(Connection connection, int tenantId) throws SQLException {

        List<String> policyIdList = new ArrayList<>();
        try (NamedPreparedStatement preparedStmt = new NamedPreparedStatement(connection,
                LIST_SUBSCRIBER_IDS_WITH_STATUS_SQL)) {
            preparedStmt.setInt(TENANT_ID, tenantId);
            try (ResultSet resultSet = preparedStmt.executeQuery()) {
                while (resultSet.next()) {
                    policyIdList.add(resultSet.getString(SUBSCRIBER_ID));
                }
            }
        }
        return policyIdList;
    }

    private void storeStatuses(Connection connection, Map<String, List<StatusHolder>> statusHolders, String about,
                               int tenantId) throws SQLException {

        String query = EntitlementConstants.Status.ABOUT_POLICY.equals(about)
                ? CREATE_POLICY_STATUS_SQL
                : CREATE_SUBSCRIBER_STATUS_SQL;

        try (NamedPreparedStatement addStatusPrepStmt = new NamedPreparedStatement(connection, query)) {
            for (Map.Entry<String, List<StatusHolder>> entry : statusHolders.entrySet()) {
                String key = entry.getKey();
                List<StatusHolder> statusHolderList = entry.getValue();

                for (StatusHolder statusHolder : statusHolderList) {
                    int version = -1;
                    if (statusHolder.getVersion() != null) {
                        version = Integer.parseInt(statusHolder.getVersion());
                    }

                    addStatusPrepStmt.setString(KEY, key);
                    addStatusPrepStmt.setString(STATUS_TYPE, statusHolder.getType());
                    addStatusPrepStmt.setBoolean(IS_SUCCESS, statusHolder.isSuccess());
                    addStatusPrepStmt.setString(USER, statusHolder.getUser());
                    addStatusPrepStmt.setString(TARGET, statusHolder.getTarget());
                    addStatusPrepStmt.setString(TARGET_ACTION, statusHolder.getTargetAction());
                    addStatusPrepStmt.setString(MESSAGE, statusHolder.getMessage());
                    addStatusPrepStmt.setTimeStamp(LOGGED_AT, new Timestamp(System.currentTimeMillis()),
                            Calendar.getInstance(TimeZone.getTimeZone(UTC)));
                    if (EntitlementConstants.Status.ABOUT_POLICY.equals(about)) {
                        addStatusPrepStmt.setInt(VERSION, version);
                    }
                    addStatusPrepStmt.setInt(TENANT_ID, tenantId);

                    addStatusPrepStmt.addBatch();
                }
            }
            addStatusPrepStmt.executeBatch();
        }
    }

    private void storePolicies(Connection connection, List<PolicyMigrationDTO> policies, int tenantId)
            throws SQLException {

        insertPolicy(connection, policies, tenantId);
        insertPolicyReferences(connection, policies, tenantId);
        if (EntitlementUtil.isPolicyMetadataStoringEnabled()) {
            insertPolicyAttributes(connection, policies, tenantId);
        }
        insertPolicyEditorData(connection, policies, tenantId);
    }

    private void insertPolicy(Connection connection, List<PolicyMigrationDTO> policies, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement createPolicyPrepStmt = new NamedPreparedStatement(connection,
                CREATE_POLICY_SQL)) {
            for (PolicyMigrationDTO migrationPolicy : policies) {
                PolicyDTO policy = migrationPolicy.getPolicyDTO();
                createPolicyPrepStmt.setString(POLICY_ID, policy.getPolicyId());
                createPolicyPrepStmt.setInt(VERSION, Integer.parseInt(policy.getVersion()));
                createPolicyPrepStmt.setBoolean(IS_IN_PDP, policy.isPromote());
                createPolicyPrepStmt.setBoolean(IS_IN_PAP, migrationPolicy.isInPAP());
                createPolicyPrepStmt.setString(POLICY, policy.getPolicy());
                createPolicyPrepStmt.setBoolean(IS_ACTIVE, policy.isActive());
                createPolicyPrepStmt.setString(POLICY_TYPE, policy.getPolicyType());
                createPolicyPrepStmt.setString(POLICY_EDITOR, policy.getPolicyEditor());
                createPolicyPrepStmt.setInt(POLICY_ORDER, policy.getPolicyOrder());
                createPolicyPrepStmt.setTimeStamp(LAST_MODIFIED_TIME, Timestamp.valueOf(policy.getLastModifiedTime()),
                        Calendar.getInstance(TimeZone.getTimeZone(UTC)));
                createPolicyPrepStmt.setString(LAST_MODIFIED_USER,
                        CarbonContext.getThreadLocalCarbonContext().getUsername());
                createPolicyPrepStmt.setInt(TENANT_ID, tenantId);

                createPolicyPrepStmt.addBatch();
            }
            createPolicyPrepStmt.executeBatch();
        }
    }

    private void insertPolicyReferences(Connection connection, List<PolicyMigrationDTO> policies, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement createPolicyReferencesPrepStmt = new NamedPreparedStatement(connection,
                CREATE_POLICY_REFS_SQL);
             NamedPreparedStatement createPolicySetReferencesPrepStmt = new NamedPreparedStatement(connection,
                     CREATE_POLICY_SET_REFS_SQL)) {

            for (PolicyMigrationDTO migrationPolicy : policies) {
                PolicyDTO policy = migrationPolicy.getPolicyDTO();
                String[] policyIdReferences = policy.getPolicyIdReferences();
                String[] policySetIdReferences = policy.getPolicySetIdReferences();
                for (String policyIdReference : policyIdReferences) {
                    createPolicyReferencesPrepStmt.setString(REFERENCE, policyIdReference);
                    createPolicyReferencesPrepStmt.setString(POLICY_ID, policy.getPolicyId());
                    createPolicyReferencesPrepStmt.setInt(VERSION, Integer.parseInt(policy.getVersion()));
                    createPolicyReferencesPrepStmt.setInt(TENANT_ID, tenantId);
                    createPolicyReferencesPrepStmt.addBatch();
                }

                for (String policySetReference : policySetIdReferences) {
                    createPolicySetReferencesPrepStmt.setString(SET_REFERENCE, policySetReference);
                    createPolicySetReferencesPrepStmt.setString(POLICY_ID, policy.getPolicyId());
                    createPolicySetReferencesPrepStmt.setInt(VERSION, Integer.parseInt(policy.getVersion()));
                    createPolicySetReferencesPrepStmt.setInt(TENANT_ID, tenantId);
                    createPolicySetReferencesPrepStmt.addBatch();
                }
            }
            createPolicyReferencesPrepStmt.executeBatch();
            createPolicySetReferencesPrepStmt.executeBatch();
        }
    }

    private void insertPolicyAttributes(Connection connection, List<PolicyMigrationDTO> policies, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement createAttributesPrepStmt = new NamedPreparedStatement(connection,
                CREATE_POLICY_ATTRIBUTES_SQL)) {
            for (PolicyMigrationDTO migrationPolicy : policies) {
                PolicyDTO policy = migrationPolicy.getPolicyDTO();
                AttributeDTO[] attributeDTOs = policy.getAttributeDTOs();
                if (attributeDTOs != null && attributeDTOs.length > 0) {
                    for (AttributeDTO attributeDTO : attributeDTOs) {
                        createAttributesPrepStmt.setString(ATTRIBUTE_ID, attributeDTO.getAttributeId());
                        createAttributesPrepStmt.setString(ATTRIBUTE_VALUE, attributeDTO.getAttributeValue());
                        createAttributesPrepStmt.setString(DATA_TYPE, attributeDTO.getAttributeDataType());
                        createAttributesPrepStmt.setString(CATEGORY, attributeDTO.getCategory());
                        createAttributesPrepStmt.setString(POLICY_ID, policy.getPolicyId());
                        createAttributesPrepStmt.setInt(VERSION, Integer.parseInt(policy.getVersion()));
                        createAttributesPrepStmt.setInt(TENANT_ID, tenantId);
                        createAttributesPrepStmt.addBatch();
                    }
                }
            }
            createAttributesPrepStmt.executeBatch();
        }
    }

    private void insertPolicyEditorData(Connection connection, List<PolicyMigrationDTO> policies, int tenantId)
            throws SQLException {

        try (NamedPreparedStatement createPolicyEditorDataPrepStmt = new NamedPreparedStatement(connection,
                CREATE_POLICY_EDITOR_DATA_SQL)) {
            for (PolicyMigrationDTO migrationPolicy : policies) {
                PolicyDTO policy = migrationPolicy.getPolicyDTO();
                String[] policyMetaData = policy.getPolicyEditorData();
                if (policyMetaData != null && policyMetaData.length > 0) {
                    int index = 0;
                    for (String policyData : policyMetaData) {
                        createPolicyEditorDataPrepStmt.setInt(EDITOR_DATA_ORDER, index);
                        createPolicyEditorDataPrepStmt.setString(EDITOR_DATA, policyData);
                        createPolicyEditorDataPrepStmt.setString(POLICY_ID, policy.getPolicyId());
                        createPolicyEditorDataPrepStmt.setInt(VERSION, Integer.parseInt(policy.getVersion()));
                        createPolicyEditorDataPrepStmt.setInt(TENANT_ID, tenantId);

                        createPolicyEditorDataPrepStmt.addBatch();
                        index++;
                    }
                }
            }
            createPolicyEditorDataPrepStmt.executeBatch();
        }
    }

    private List<String> getExistingPolicyIds(Connection connection, int tenantId) throws SQLException {

        List<String> policyIds = new ArrayList<>();
        try (NamedPreparedStatement prepStmt = new NamedPreparedStatement(connection, LIST_POLICY_IDS_SQL)) {
            prepStmt.setInt(TENANT_ID, tenantId);

            try (ResultSet resultSet = prepStmt.executeQuery()) {
                while (resultSet.next()) {
                    policyIds.add(resultSet.getString(POLICY_ID));
                }
                if (policyIds.isEmpty()) {
                    LOG.debug("No PAP policies found");
                }
                return policyIds;
            }
        }
    }
}
