/*
 * Copyright (c) 2017, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
 *
 * Licensed 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;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.wso2.carbon.database.utils.jdbc.JdbcTemplate;
import org.wso2.carbon.database.utils.jdbc.exceptions.DataAccessException;
import org.wso2.carbon.database.utils.jdbc.exceptions.TransactionException;
import org.wso2.carbon.identity.core.migrate.MigrationClient;
import org.wso2.carbon.identity.core.migrate.MigrationClientException;
import org.wso2.carbon.identity.core.util.JdbcUtils;
import org.wso2.carbon.is.migration.config.Config;
import org.wso2.carbon.is.migration.util.Constant;
import org.wso2.carbon.is.migration.util.Utility;

import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Objects;
import java.util.UUID;

import static org.wso2.carbon.is.migration.util.Constant.COMPLETED;
import static org.wso2.carbon.is.migration.util.Constant.EXPIRES_AT_COLUMN;
import static org.wso2.carbon.is.migration.util.Constant.FAILED;
import static org.wso2.carbon.is.migration.util.Constant.IN_PROGRESS;
import static org.wso2.carbon.is.migration.util.Constant.LOCK_TIMEOUT_MINUTES;
import static org.wso2.carbon.is.migration.util.Constant.LOCKED_BY_COLUMN;
import static org.wso2.carbon.is.migration.util.Constant.MAX_POLL_ATTEMPTS;
import static org.wso2.carbon.is.migration.util.Constant.DB_LOCK_MIN_SUPPORTED_VERSION;
import static org.wso2.carbon.is.migration.util.Constant.NOT_STARTED;
import static org.wso2.carbon.is.migration.util.Constant.POLLING_INTERVAL_MS;
import static org.wso2.carbon.is.migration.util.Constant.STATUS_COLUMN;
import static org.wso2.carbon.is.migration.util.SQLConstants.RELEASE_LOCK_AND_UPDATE_STATUS_QUERY;
import static org.wso2.carbon.is.migration.util.SQLConstants.SELECT_STATUS_QUERY;
import static org.wso2.carbon.is.migration.util.SQLConstants.UPDATE_STATUS_QUERY;

/**
 * MigrationClientImpl is the one that trigger by the relevant component trigger to start the migration service.
 */
public class MigrationClientImpl implements MigrationClient {

    private static final Logger log = LoggerFactory.getLogger(MigrationClientImpl.class);
    boolean isSuccessful = false;
    private String destinationVersion;
    private final String LOCKED_BY = UUID.randomUUID().toString();
    private String migrationState = NOT_STARTED;

    @Override
    public void execute() throws MigrationClientException {

        boolean isLegacyMigration;

        Config config = Config.getInstance();

        if (!config.isMigrationEnable()) {
            return;
        }
        destinationVersion = config.getMigrateVersion();
        isLegacyMigration = isLegacyMigration(config);

        handleMigrationProcess(config, isLegacyMigration);
    }

    private boolean isLegacyMigration(Config config) {

        return DB_LOCK_MIN_SUPPORTED_VERSION.compareTo(config.getCurrentVersion()) > 0 || config.isLegacyMigration();
    }

    private void releaseLock(boolean needsLockRelease) throws MigrationClientException {

        if (needsLockRelease) {
            try {
                releaseLock();
            } catch (DataAccessException e) {
                log.info("Failed to release the lock.");
                throw new MigrationClientException(e.getMessage(), e);
            }
        } else {
            log.info("Lock release is not needed.");
        }
    }

    private void handleMigrationProcess(Config config, boolean isLegacyMigration) throws MigrationClientException {

        boolean needsLockRelease = true;
        try {
            if (isLegacyMigration) {
                log.info("This is a legacy migration hence executing the client without database lock concept.");
                executeMigration(config);
                needsLockRelease = false;
                return;
            }

            boolean isForceRerun = config.isForceRerun();
            boolean performUpgrade = acquireLock(isForceRerun) || shouldRetryMigration(isForceRerun);

            if (performUpgrade) {
                log.info("Migration needs to be done hence executing the migration.");
                executeMigration(config);
                migrationState = COMPLETED;
            } else {
                needsLockRelease = false;
                log.info("Migration is not needed. Hence continue the start up.");
            }
        } catch (DataAccessException e) {
            migrationState = FAILED;
            throw new MigrationClientException(e.getMessage(), e);
        } finally {
            releaseLock(needsLockRelease);
        }
    }

    private boolean shouldRetryMigration(boolean isForceRerun) throws DataAccessException, MigrationClientException {

        // If the migration is already completed and isForceRerun is false, then no need to wait for
        // the migration completion.
        if (COMPLETED.equalsIgnoreCase(migrationState) & !isForceRerun) {
            log.info("Migration is already completed and it not a force re run hence retry migration is not needed.");
            return false;
        }
        waitForMigrationCompletion(destinationVersion);

        return (!COMPLETED.equalsIgnoreCase(migrationState) || isForceRerun) && acquireLock(isForceRerun);
    }

    private void executeMigration(Config config) {

        try {
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");
            log.info("............................................................................................");

            VersionMigrationHolder versionMigrationHolder = VersionMigrationHolder.getInstance();
            List<VersionMigration> versionMigrationList = versionMigrationHolder.getVersionMigrationList();

            log.info("Migration Versions List.........................");
            for (VersionMigration versionMigration : versionMigrationList) {
                log.info(versionMigration.getPreviousVersion() + " to " + versionMigration.getCurrentVersion());
            }

            boolean isMigrationStarted = false;
            if (Utility.isMigrateTenantRange()) {
                log.info("Migration started for the tenant range " + Utility.getMigrationStartingTenantID() + " - "
                        + Utility.getMigrationEndingTenantID());
            }

            for (VersionMigration versionMigration : versionMigrationList) {
                log.info(Constant.MIGRATION_LOG + "Start Version : " + versionMigration.getPreviousVersion() + " to "
                        + versionMigration.getCurrentVersion());
                if (!isMigrationStarted && versionMigration.getPreviousVersion().equals(config.getCurrentVersion())) {
                    versionMigration.migrate();
                    isMigrationStarted = true;
                    if (versionMigration.getCurrentVersion().equals(config.getMigrateVersion())) {
                        break;
                    }
                    continue;
                }

                if (isMigrationStarted) {
                    versionMigration.migrate();
                    if (versionMigration.getCurrentVersion().equals(config.getMigrateVersion())) {
                        break;
                    }
                }
            }
            log.info(Constant.MIGRATION_LOG + "Execution was done through all the requested version list without "
                    + "having unexpected issues. There may be some steps that is not executed correctly but bypass "
                    + "that because of enabling 'continueOnError' property. Please see the above logs to more "
                    + "details.");
            isSuccessful = true;
        } catch (Throwable e) {
            log.error("Migration process was stopped.", e);
        }

        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
        log.info("............................................................................................");
    }

    @Override
    public void executeAfterStartup() {

        Config config = Config.getInstance();

        if (config.isMigrationEnable() && config.isExitOnComplete()) {
            log.info("Shutting down as the migration execution is completed.");
            if (isSuccessful) {
                System.exit(0);
            } else {
                System.exit(1);
            }
        }
    }

    private boolean acquireLock(boolean isForceRerun) throws MigrationClientException {

        try {
            JdbcTemplate jdbcTemplate = JdbcUtils.getNewTemplate();
            return jdbcTemplate.withTransaction(template -> {
                if (isLockAvailable(isForceRerun)) {
                    log.info("Lock is available to allocate.");
                    return allocateLock();
                }
                return false;
            });
        } catch (TransactionException e) {
            migrationState = FAILED;
            throw new MigrationClientException("Error while aquiring the lock.", e);
        }
    }

    private boolean isLockAvailable(boolean isForceRerun) throws TransactionException {

        JdbcTemplate jdbcTemplate = JdbcUtils.getNewTemplate();
        return jdbcTemplate.withTransaction((template) -> template.fetchSingleRecord(SELECT_STATUS_QUERY, (resultSet, rowNumber) -> {
            migrationState = resultSet.getString(STATUS_COLUMN);
            Timestamp expiresAt = resultSet.getTimestamp(EXPIRES_AT_COLUMN);
            LocalDateTime now = LocalDateTime.now();
            Timestamp currentTimeStamp = Timestamp.valueOf(now);
            if (COMPLETED.equalsIgnoreCase(migrationState)) {
                return isForceRerun;
            } else {
                return StringUtils.isBlank(resultSet.getString(LOCKED_BY_COLUMN)) || expiresAt.compareTo(currentTimeStamp) < 0 ||
                        FAILED.equalsIgnoreCase(migrationState);
            }
        }, preparedStatement -> {
            preparedStatement.setString(1, destinationVersion);
        }));
    }

    private boolean allocateLock() throws TransactionException {

        LocalDateTime now = LocalDateTime.now();
        Timestamp currentTimestamp = Timestamp.valueOf(now);
        Timestamp expiredTimestamp = Timestamp.valueOf(now.plusMinutes(LOCK_TIMEOUT_MINUTES));
        // Allocate the lock by updating the row
        JdbcTemplate jdbcTemplate = JdbcUtils.getNewTemplate();
        int affectedRows = jdbcTemplate.withTransaction((template) -> template.
                executeUpdateWithAffectedRows(UPDATE_STATUS_QUERY, preparedStatement -> {
                    preparedStatement.setString(1, LOCKED_BY);
                    preparedStatement.setString(2, IN_PROGRESS);
                    preparedStatement.setTimestamp(3, currentTimestamp);
                    preparedStatement.setTimestamp(4, expiredTimestamp);
                    preparedStatement.setString(5, destinationVersion);
                }));
        if (affectedRows > 0) {
            log.info("Lock allocated to the node: {}", LOCKED_BY);
        } else{
            log.info("Lock could not be allocated to the node: {}", LOCKED_BY);
        }
        return affectedRows > 0;
    }

    public void waitForMigrationCompletion(String isVersion) throws DataAccessException {

        int pollCount = 0;
        while (pollCount < MAX_POLL_ATTEMPTS) {
            String lockStatus = checkLockStatus(isVersion);
            switch (Objects.requireNonNull(lockStatus)) {
                case IN_PROGRESS:
                    log.info("Migration in progress. Retrying...");
                    sleepBeforeRetry();
                    break;

                case COMPLETED:
                    migrationState = COMPLETED;
                    log.info("Migration completed. Proceeding with startup.");
                    // Exit the loop and continue server startup
                    return;

                case FAILED:
                    log.info("Migration failed. Allowing another node to retry.");
                    return;

                default:
                    log.info("Unexpected lock status: {}", lockStatus);
                    return;
            }
            pollCount++;
        }
        log.info("Exceeded maximum polling attempts. Exiting.");
    }

    private String checkLockStatus(String isVersion) throws DataAccessException {

        JdbcTemplate jdbcTemplate = JdbcUtils.getNewTemplate();
        List<String> migrationStatusList = jdbcTemplate.executeQuery(SELECT_STATUS_QUERY,
                (resultSet, rowNumber) -> resultSet.getString(STATUS_COLUMN),
                preparedStatement -> preparedStatement.setString(1, isVersion));

        return CollectionUtils.isNotEmpty(migrationStatusList) ? migrationStatusList.get(0) : null;
    }

    private void sleepBeforeRetry() {

        try {
            // Wait before retrying.
            Thread.sleep(POLLING_INTERVAL_MS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.error("Polling interrupted", e);
        }
    }

    private void releaseLock() throws DataAccessException {

        JdbcTemplate jdbcTemplate = JdbcUtils.getNewTemplate();
        jdbcTemplate.executeUpdate(RELEASE_LOCK_AND_UPDATE_STATUS_QUERY, preparedStatement -> {
            preparedStatement.setString(1, migrationState);
            preparedStatement.setString(2, destinationVersion);
            preparedStatement.setString(3, LOCKED_BY);
        });

        log.info("Lock released by node: {}", LOCKED_BY);
    }
}
