/*
 * Copyright (c) 2025, 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.wso2.carbon.database.utils.jdbc.NamedJdbcTemplate;
import org.wso2.carbon.database.utils.jdbc.NamedPreparedStatement;
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.MigrationClientException;
import org.wso2.carbon.identity.core.model.SAMLSSOServiceProviderDO;
import org.wso2.carbon.identity.core.model.ServiceProviderProperty;
import org.wso2.carbon.identity.core.util.JdbcUtils;

import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.Calendar;
import java.util.Date;
import java.util.List;
import java.util.TimeZone;
import org.wso2.carbon.identity.sso.saml.SAMLSSOConstants;

import static org.wso2.carbon.is.migration.service.v710.constants.SAMLSQLConstants.ADD_SAML2_SSO_CONFIG;
import static org.wso2.carbon.is.migration.service.v710.constants.SAMLSQLConstants.ADD_SAML_SSO_ATTR;
import static org.wso2.carbon.is.migration.service.v710.constants.SAMLSQLConstants.GET_SAML_SP_ID_BY_ISSUER;
import static org.wso2.carbon.is.migration.service.v710.constants.SAMLSQLConstants.SAMLTableColumns.*;

import static java.time.ZoneOffset.UTC;

/**
 * This class handles the migrations related to SAML.
 */
public class SAMLDAO {

    private static final Logger LOG = LoggerFactory.getLogger(SAMLDAO.class);
    private final Calendar CALENDAR = Calendar.getInstance(TimeZone.getTimeZone(UTC));
    public static final String SAML_SCHEMA_VERSION = "1.0.0";

    public void migrateSAMLServiceProviders(List<SAMLSSOServiceProviderDO> samlServiceProviders, boolean isDryRun,
                                            int tenantId) throws MigrationClientException {

        if (samlServiceProviders == null || samlServiceProviders.isEmpty()) {
            LOG.info("DM-SAML>> No SAML service providers to migrate in tenant: {}.", tenantId);
            return;
        }
        LOG.info("DM-SAML>> {} SAML service providers picked for migration from Registry in tenant: {}.",
                samlServiceProviders.size(), tenantId);
        // Perform the migration in a transaction
        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();
        try {
            // Add the SAML service providers to the DB
            namedJdbcTemplate.withTransaction(template -> {
                if (!isDryRun) {
                    storeServiceProviders(samlServiceProviders, tenantId);
                } else {
                    LOG.info("DM-SAML>> Dry run enabled. Skipping the SAML service provider migration in tenant: {}.",
                            tenantId);
                }
                return null;
            });
        } catch (TransactionException e) {
            throw new MigrationClientException(
                    String.format("Error while migrating service providers in tenant: %s.", tenantId), e);
        }
    }

    private void storeServiceProviders(List<SAMLSSOServiceProviderDO> samlServiceProviders, int tenantId)
            throws DataAccessException, TransactionException {

        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();
        int serviceProviderErrorCount = 0;
        for (SAMLSSOServiceProviderDO serviceProviderDO : samlServiceProviders) {
            if (!isServiceProviderExists(serviceProviderDO.getIssuer(), tenantId)) {
                namedJdbcTemplate.withTransaction(template -> {
                    processAddServiceProvider(serviceProviderDO, tenantId);
                    processAddSPProperties(serviceProviderDO, tenantId);
                    return null;
                });
            } else {
                serviceProviderErrorCount++;
                LOG.warn(
                        "DM-SAML>> SAML service provider with issuer: {} already exists in tenant: {}. Skipping service provider. ",
                        serviceProviderDO.getIssuer(), tenantId);
            }
        }
        if (serviceProviderErrorCount > 0) {
            LOG.warn("DM-SAML>> {} SAML service providers already exist in tenant: {}.", serviceProviderErrorCount,
                    tenantId);
        }
    }

    private boolean isServiceProviderExists(String issuer, int tenantId) throws DataAccessException {

        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();
        Integer serviceProviderId = namedJdbcTemplate.fetchSingleRecord(GET_SAML_SP_ID_BY_ISSUER,
                (resultSet, rowNumber) -> resultSet.getInt(ID), preparedStatement -> {
                    preparedStatement.setString(ISSUER, issuer);
                    preparedStatement.setInt(TENANT_ID, tenantId);
                });
        return serviceProviderId != null;

    }

    private void processAddServiceProvider(SAMLSSOServiceProviderDO serviceProviderDO, int tenantId)
            throws DataAccessException {

        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();
        namedJdbcTemplate.executeInsert(ADD_SAML2_SSO_CONFIG,
                namedPreparedStatement -> setServiceProviderParameters(namedPreparedStatement, serviceProviderDO,
                        tenantId), serviceProviderDO, false);
    }

    private void processAddSPProperties(SAMLSSOServiceProviderDO serviceProviderDO, int tenantId)
            throws DataAccessException {

        List<ServiceProviderProperty> properties = serviceProviderDO.getMultiValuedProperties();
        int serviceProviderId = processGetServiceProviderId(serviceProviderDO.getIssuer(), tenantId);

        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();

        namedJdbcTemplate.executeBatchInsert(ADD_SAML_SSO_ATTR, (namedPreparedStatement -> {
            for (ServiceProviderProperty property : properties) {
                namedPreparedStatement.setInt(SP_ID, serviceProviderId);
                namedPreparedStatement.setString(PROPERTY_NAME, property.getKey());
                namedPreparedStatement.setString(PROPERTY_VALUE, property.getValue());
                namedPreparedStatement.addBatch();
            }
        }), serviceProviderDO);
    }

    private int processGetServiceProviderId(String issuer, int tenantId) throws DataAccessException {

        NamedJdbcTemplate namedJdbcTemplate = JdbcUtils.getNewNamedJdbcTemplate();
        Integer serviceProviderId = namedJdbcTemplate.fetchSingleRecord(GET_SAML_SP_ID_BY_ISSUER,
                (resultSet, rowNumber) -> resultSet.getInt(ID), namedPreparedStatement -> {
                    namedPreparedStatement.setString(ISSUER, issuer);
                    namedPreparedStatement.setInt(TENANT_ID, tenantId);
                });
        if (serviceProviderId == null) {
            throw new DataAccessException("No record found for the given issuer: " + issuer);
        }
        return serviceProviderId.intValue();
    }

    private void setServiceProviderParameters(NamedPreparedStatement statement,
                                              SAMLSSOServiceProviderDO serviceProviderDO, int tenantId)
            throws SQLException {

        Timestamp currentTime = new Timestamp(new Date().getTime());
        statement.setInt(TENANT_ID, tenantId);
        statement.setString(ISSUER, serviceProviderDO.getIssuer());
        statement.setString(DEFAULT_ASSERTION_CONSUMER_URL, serviceProviderDO.getDefaultAssertionConsumerUrl());
        statement.setString(NAME_ID_FORMAT, serviceProviderDO.getNameIDFormat());
        statement.setString(CERT_ALIAS, serviceProviderDO.getCertAlias());
        statement.setBoolean(REQ_SIG_VALIDATION, serviceProviderDO.isDoValidateSignatureInRequests());
        statement.setBoolean(SIGN_RESPONSE, serviceProviderDO.isDoSignResponse());
        statement.setBoolean(SIGN_ASSERTIONS, serviceProviderDO.isDoSignAssertions());
        statement.setString(SIGNING_ALGO, serviceProviderDO.getSigningAlgorithmUri());
        statement.setString(DIGEST_ALGO, serviceProviderDO.getDigestAlgorithmUri());
        statement.setBoolean(ENCRYPT_ASSERTION, serviceProviderDO.isDoEnableEncryptedAssertion());
        statement.setString(ASSERTION_ENCRYPTION_ALGO, serviceProviderDO.getAssertionEncryptionAlgorithmUri());
        statement.setString(KEY_ENCRYPTION_ALGO, serviceProviderDO.getKeyEncryptionAlgorithmUri());
        statement.setBoolean(ATTR_PROFILE_ENABLED, serviceProviderDO.isEnableAttributesByDefault());
        statement.setString(ATTR_SERVICE_INDEX, serviceProviderDO.getAttributeConsumingServiceIndex());
        statement.setBoolean(SLO_PROFILE_ENABLED, serviceProviderDO.isDoSingleLogout());
        statement.setString(SLO_METHOD, serviceProviderDO.getSingleLogoutMethod());
        statement.setString(SLO_RESPONSE_URL, serviceProviderDO.getSloResponseURL());
        statement.setString(SLO_REQUEST_URL, serviceProviderDO.getSloRequestURL());
        statement.setBoolean(IDP_INIT_SSO_ENABLED, serviceProviderDO.isIdPInitSSOEnabled());
        statement.setBoolean(IDP_INIT_SLO_ENABLED, serviceProviderDO.isIdPInitSLOEnabled());
        statement.setBoolean(QUERY_REQUEST_PROFILE_ENABLED, serviceProviderDO.isAssertionQueryRequestProfileEnabled());
        statement.setBoolean(ECP_ENABLED, serviceProviderDO.isSamlECP());
        statement.setBoolean(ARTIFACT_BINDING_ENABLED, serviceProviderDO.isEnableSAML2ArtifactBinding());
        statement.setBoolean(ARTIFACT_RESOLVE_REQ_SIG_VALIDATION,
                serviceProviderDO.isDoValidateSignatureInArtifactResolve());
        statement.setString(IDP_ENTITY_ID_ALIAS, serviceProviderDO.getIdpEntityIDAlias());
        statement.setString(ISSUER_QUALIFIER, serviceProviderDO.getIssuerQualifier());
        statement.setString(SUPPORTED_ASSERTION_QUERY_REQUEST_TYPES,
                serviceProviderDO.getSupportedAssertionQueryRequestTypes());
        statement.setString(VERSION, SAML_SCHEMA_VERSION);
        statement.setTimeStamp(CREATED_AT, currentTime, CALENDAR);
        statement.setTimeStamp(UPDATED_AT, currentTime, CALENDAR);
        if (serviceProviderDO.getAttributeNameFormat() != null) {
            statement.setString(ATTR_NAME_FORMAT, serviceProviderDO.getAttributeNameFormat());
        } else {
            statement.setString(ATTR_NAME_FORMAT, SAMLSSOConstants.NameFormat.BASIC.toString());
        }
    }
}
