/*
 * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com). All Rights Reserved.
 *
 * This software is the property of WSO2 LLC. and its suppliers, if any.
 * Dissemination of any information or reproduction of any material contained
 * herein in any form is strictly forbidden, unless permitted by WSO2 expressly.
 * You may not alter or remove any copyright or other notice from copies of this content.
 */

package org.wso2.carbon.is.migration.service.v700.util;

import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.impl.builder.StAXOMBuilder;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.oltu.oauth2.common.message.types.GrantType;
import org.apache.oltu.oauth2.common.message.types.ResponseType;
import org.wso2.carbon.context.CarbonContext;
import org.wso2.carbon.database.utils.jdbc.NamedJdbcTemplate;
import org.wso2.carbon.identity.core.migrate.MigrationClientException;
import org.wso2.carbon.identity.core.util.IdentityCoreConstants;
import org.wso2.carbon.identity.oauth.common.OAuthConstants;
import org.wso2.carbon.user.api.UserStoreException;
import org.wso2.carbon.user.core.util.DatabaseUtil;
import org.wso2.carbon.utils.CarbonUtils;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;

import javax.sql.DataSource;
import javax.xml.namespace.QName;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;

import static org.wso2.carbon.identity.oauth2.util.OAuth2Util.OPENID_CONNECT;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.APPLICATION_ACCESS_TOKEN_VALIDATION_PERIOD;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.CONFIG_ELEM_OAUTH;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.GRANT_TYPE_NAME;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.GRANT_TYPE_ORGANIZATION_SWITCH;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.IMPLICIT;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.OPENID_CONNECT_IDTOKEN_EXPIRATION;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.REFRESH_TOKEN_VALIDITY_PERIOD;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.RESPONSE_TYPE_ID_TOKEN;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.RESPONSE_TYPE_TOKEN;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.RESP_TYPE_NAME;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.SUPPORTED_GRANT_TYPE;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.SUPPORTED_GRANT_TYPES;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.SUPPORTED_RESP_TYPE;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.SUPPORTED_RESP_TYPES;
import static org.wso2.carbon.is.migration.service.v700.constant.MigratorConstants.USER_ACCESS_TOKEN_DEFAULT_VALIDITY_PERIOD;

public class Utils {

    private static OMElement oAuthElem;
    private static long userAccessTokenValidityPeriodInSeconds = 3600;
    private static long applicationAccessTokenValidityPeriodInSeconds = 3600;
    private static long refreshTokenValidityPeriodInSeconds = 24L * 3600;
    private static long openIDConnectIDTokenExpiryTimeInSeconds = 3600;

    private static final Log LOG = LogFactory.getLog(Utils.class);
    private static Map<String, String> oAuthScopeBindings = null;
    private static DataSource umDataSource = null;

    private static final List<String> supportedGrantTypes = new ArrayList<>();
    private static final List<String> supportedResponseTypes = new ArrayList<>();

    public static NamedJdbcTemplate getNewTemplateForSharedDB() throws MigrationClientException {

        getUMDataSource();
        return new NamedJdbcTemplate(umDataSource);
    }

    public static DataSource getUMDataSource() throws MigrationClientException {

        if (umDataSource == null) {
            try {
                umDataSource = DatabaseUtil.getRealmDataSource(CarbonContext.getThreadLocalCarbonContext().
                        getUserRealm().getRealmConfiguration());
            } catch (UserStoreException e) {
                throw new MigrationClientException("Error while retrieving realm data source.", e);
            }
        }
        return umDataSource;
    }

    public static Map<String, String> loadOauthScopeBinding() {

        if (oAuthScopeBindings != null) {
            return oAuthScopeBindings;
        }
        String configDirPath = CarbonUtils.getCarbonConfigDirPath();
        String confXml = Paths.get(configDirPath, "identity", OAuthConstants.OAUTH_SCOPE_BINDING_PATH).toString();
        File configFile = new File(confXml);
        if (!configFile.exists()) {
            LOG.warn("OAuth scope binding File is not present at: " + confXml);
            return null;
        }

        XMLStreamReader parser = null;
        try (InputStream stream = Files.newInputStream(configFile.toPath())) {

            parser = XMLInputFactory.newInstance()
                    .createXMLStreamReader(stream);
            StAXOMBuilder builder = new StAXOMBuilder(parser);
            OMElement documentElement = builder.getDocumentElement();
            Iterator iterator = documentElement.getChildElements();
            Map<String, String> scopeBindings = new HashMap<>();
            while (iterator.hasNext()) {
                OMElement omElement = (OMElement) iterator.next();
                String scopeName = omElement.getAttributeValue(new QName("name"));
                List<String> bindingPermissions = loadScopePermissions(omElement);
                if (bindingPermissions.size() == 0) {
                    continue;
                }
                scopeBindings.put(bindingPermissions.get(0), scopeName);
                if (bindingPermissions.size() > 1) {
                    throw new RuntimeException("Scope binding can only have one permission.");
                }
            }
            oAuthScopeBindings = scopeBindings;
            return scopeBindings;
        } catch (XMLStreamException e) {
            LOG.warn("Error while streaming oauth-scope-bindings config.", e);
        } catch (IOException e) {
            LOG.warn("Error while loading oauth-scope-bindings config.", e);
        } finally {
            try {
                if (parser != null) {
                    parser.close();
                }
            } catch (XMLStreamException e) {
                LOG.error("Error while closing XML stream", e);
            }
        }
        return null;
    }

    private static List<String> loadScopePermissions(OMElement configElement) {

        List<String> permissions = new ArrayList<>();
        Iterator it = configElement.getChildElements();
        while (it.hasNext()) {
            OMElement element = (OMElement) it.next();
            Iterator permissionIterator = element.getChildElements();
            while (permissionIterator.hasNext()) {
                OMElement permissionElement = (OMElement) permissionIterator.next();
                if ("Permission".equals(permissionElement.getLocalName())) {
                    String permission = permissionElement.getText();
                    permissions.add(permission);
                }
            }
        }
        return permissions;
    }

    /**
     * Get Allowed grant types configured.
     *
     * @return List of Grant Types.
     * @throws MigrationClientException If an error occurs while reading the identity.xml file.
     */
    public static List<String> getAllowedGrantTypes() throws MigrationClientException {

        List<String> allowedGrantSet = getSupportedGrantTypes();
        Set<String> modifiableGrantSet = new HashSet<>(Objects.requireNonNull(allowedGrantSet));
        if (isImplicitGrantEnabled()) {
            modifiableGrantSet.add(IMPLICIT);
        }

        return new ArrayList<>(modifiableGrantSet);
    }

    /**
     * Get configured ID token validity period time in seconds.
     *
     * @return ID token token validity period configured.
     */
    public static long getOpenIDConnectIDTokenValidityPeriodInSeconds() {

        return openIDConnectIDTokenExpiryTimeInSeconds;
    }

    /**
     * Get configured user access token expiration time in seconds.
     *
     * @return User access token validity period configured.
     */
    public static long getUserAccessTokenValidityPeriodInSeconds() {

        return userAccessTokenValidityPeriodInSeconds;
    }

    /**
     * Get configured refresh token expiration time in seconds.
     *
     * @return Refresh token validity period configured.
     */
    public static long getRefreshTokenValidityPeriodInSeconds() {

        return refreshTokenValidityPeriodInSeconds;
    }

    /**
     * Get configured application access token expiration time in seconds.
     *
     * @return Application access token validity period configured.
     */
    public static long getApplicationAccessTokenValidityPeriodInSeconds() {

        return applicationAccessTokenValidityPeriodInSeconds;
    }

    /**
     * Reads configured default validity periods for tokens.
     */
    public static void parseDefaultValidityPeriods() throws MigrationClientException {

        parseOAuthOMElement();
        if (oAuthElem == null) {
            LOG.error("OAuth element is not available in identity.xml");
            throw new MigrationClientException("Error while reading identity.xml");
        }

        // Set the access token default timeout.
        OMElement accessTokenTimeoutElem = oAuthElem.getFirstChildWithName(getQNameWithIdentityNS(
                USER_ACCESS_TOKEN_DEFAULT_VALIDITY_PERIOD));
        if (accessTokenTimeoutElem != null) {
            userAccessTokenValidityPeriodInSeconds = Long.parseLong(accessTokenTimeoutElem.getText().trim());
        }

        // Set the application access token default timeout.
        OMElement applicationAccessTokTimeoutElem = oAuthElem.getFirstChildWithName(
                getQNameWithIdentityNS(
                        APPLICATION_ACCESS_TOKEN_VALIDATION_PERIOD));
        if (applicationAccessTokTimeoutElem != null) {
            applicationAccessTokenValidityPeriodInSeconds = Long.parseLong(
                    applicationAccessTokTimeoutElem.getText().trim());
        }

        // Set the refresh token default timeout.
        OMElement refreshTokenTimeoutElem = oAuthElem.getFirstChildWithName(
                getQNameWithIdentityNS(REFRESH_TOKEN_VALIDITY_PERIOD));
        if (refreshTokenTimeoutElem != null) {
            refreshTokenValidityPeriodInSeconds = Long.parseLong(refreshTokenTimeoutElem.getText().trim());
        }

        // Set the id token default timeout.
        parseOpenIDConnectIDTokenExpirationTime(oAuthElem);
    }

    /**
     * Convert String timestamp to sql timestamp.
     *
     * @param timestamp string timestamp.
     * @return sql timestamp.
     */
    public static Timestamp formatTimestamp(String timestamp) {

        String formattedTimestamp = timestamp.replace("T", " ").replace("Z", "");
        return Timestamp.valueOf(formattedTimestamp);
    }

    private static void parseOpenIDConnectIDTokenExpirationTime(OMElement oauthConfigElem) {

        OMElement openIDConnectConfigElem =
                oauthConfigElem.getFirstChildWithName(getQNameWithIdentityNS(OPENID_CONNECT));

        if (openIDConnectConfigElem != null) {
            OMElement openIDConnectIDTokenExpirationElem = openIDConnectConfigElem.getFirstChildWithName(
                    getQNameWithIdentityNS(OPENID_CONNECT_IDTOKEN_EXPIRATION));
            if (openIDConnectIDTokenExpirationElem != null) {
                openIDConnectIDTokenExpiryTimeInSeconds = Long.parseLong(
                        openIDConnectIDTokenExpirationElem.getText().trim());
            }
        }
    }

    private static List<String> getSupportedGrantTypes() throws MigrationClientException {

        parseOAuthOMElement();

        if (oAuthElem == null) {
            LOG.error("OAuth element is not available in identity.xml");
            return null;
        }
        parseSupportedGrantTypesConfig(oAuthElem);
        parseSupportedResponseTypesConfig(oAuthElem);

        return supportedGrantTypes;
    }

    private static boolean isImplicitGrantEnabled() {

        List<String> responseTypeHandlers = getSupportedResponseTypes();
        for (String responseType : responseTypeHandlers) {
            if (responseType.contains(RESPONSE_TYPE_TOKEN) || responseType.contains(RESPONSE_TYPE_ID_TOKEN)) {
                return true;
            }
        }
        return false;
    }

    private static void parseSupportedGrantTypesConfig(OMElement oauthConfigElem) {

        if (supportedGrantTypes.isEmpty()) {
            OMElement supportedGrantTypesElem =
                    oauthConfigElem.getFirstChildWithName(getQNameWithIdentityNS(SUPPORTED_GRANT_TYPES));

            if (supportedGrantTypesElem != null) {
                Iterator<OMElement> iterator = supportedGrantTypesElem
                        .getChildrenWithName(getQNameWithIdentityNS(SUPPORTED_GRANT_TYPE));
                while (iterator.hasNext()) {
                    OMElement supportedGrantTypeElement = iterator.next();
                    OMElement grantTypeNameElement = supportedGrantTypeElement
                            .getFirstChildWithName(getQNameWithIdentityNS(GRANT_TYPE_NAME));
                    String grantTypeName = null;
                    if (grantTypeNameElement != null) {
                        grantTypeName = grantTypeNameElement.getText();
                    }

                    if (StringUtils.isNotEmpty(grantTypeName)) {
                        supportedGrantTypes.add(grantTypeName);
                    }
                }
            } else {
                // if this element is not present, assume the default case.
                LOG.warn("'SupportedGrantTypes' element not configured in identity.xml. " +
                        "Therefore instantiating default grant type handlers");

                supportedGrantTypes.add(GrantType.AUTHORIZATION_CODE.toString());
                supportedGrantTypes.add(GrantType.CLIENT_CREDENTIALS.toString());
                supportedGrantTypes.add(GrantType.PASSWORD.toString());
                supportedGrantTypes.add(GrantType.REFRESH_TOKEN.toString());
                supportedGrantTypes.add(org.wso2.carbon.identity.oauth.common.GrantType.SAML20_BEARER.toString());
                supportedGrantTypes.add(org.wso2.carbon.identity.oauth.common.GrantType.IWA_NTLM.toString());
                supportedGrantTypes.add(GRANT_TYPE_ORGANIZATION_SWITCH);
            }
        }

        if (LOG.isDebugEnabled()) {
            for (String grantTypeName : supportedGrantTypes) {
                LOG.debug(grantTypeName + " added.");
            }
        }
    }

    private static List<String> getSupportedResponseTypes() {

        return supportedResponseTypes;
    }

    private static void parseSupportedResponseTypesConfig(OMElement oauthConfigElem) {

        if (supportedResponseTypes.isEmpty()) {
            OMElement supportedRespTypesElem =
                    oauthConfigElem.getFirstChildWithName(getQNameWithIdentityNS(SUPPORTED_RESP_TYPES));

            if (supportedRespTypesElem != null) {
                Iterator<OMElement> iterator = supportedRespTypesElem
                        .getChildrenWithName(getQNameWithIdentityNS(SUPPORTED_RESP_TYPE));
                while (iterator.hasNext()) {
                    OMElement supportedResponseTypeElement = iterator.next();
                    OMElement responseTypeNameElement = supportedResponseTypeElement.
                            getFirstChildWithName(
                                    getQNameWithIdentityNS(RESP_TYPE_NAME));
                    String responseTypeName = null;
                    if (responseTypeNameElement != null) {
                        responseTypeName = responseTypeNameElement.getText();
                    }
                    if (responseTypeName != null && !responseTypeName.isEmpty()) {
                        supportedResponseTypes.add(responseTypeName);
                    }
                }
            } else {
                // if this element is not present, assume the default case.
                LOG.warn("'SupportedResponseTypes' element not configured in identity.xml. " +
                        "Therefore instantiating default response type handlers");
                supportedResponseTypes.add(ResponseType.CODE.toString());
                supportedResponseTypes.add(ResponseType.TOKEN.toString());
                supportedResponseTypes.add(OAuthConstants.ID_TOKEN);
                supportedResponseTypes.add(OAuthConstants.IDTOKEN_TOKEN);
                supportedResponseTypes.add(OAuthConstants.CODE_TOKEN);
                supportedResponseTypes.add(OAuthConstants.CODE_IDTOKEN);
                supportedResponseTypes.add(OAuthConstants.CODE_IDTOKEN_TOKEN);
            }
        }

        if (LOG.isDebugEnabled()) {
            for (String responseTypeName : supportedResponseTypes) {
                LOG.debug(responseTypeName + " added.");
            }
        }
    }

    private static void parseOAuthOMElement() throws MigrationClientException {

        if (oAuthElem == null) {
            String configDirPath = CarbonUtils.getCarbonConfigDirPath();
            String filePath = Paths.get(configDirPath, "identity", "identity.xml").toString();
            File identityConfigXml = new File(filePath);
            InputStream inStream = null;

            try {
                if (identityConfigXml.exists()) {
                    inStream = new FileInputStream(identityConfigXml);
                }
                if (inStream == null) {
                    String message = "Identity configuration not found at: " + configDirPath;
                    if (LOG.isDebugEnabled()) {
                        LOG.debug(message);
                    }
                    throw new FileNotFoundException(message);
                }
                OMElement rootElement = new StAXOMBuilder(inStream).getDocumentElement();

                if (rootElement == null) {
                    throw new MigrationClientException("Error while reading identity.xml");
                }

                oAuthElem = rootElement.getFirstChildWithName(getQNameWithIdentityNS(CONFIG_ELEM_OAUTH));

            } catch (FileNotFoundException | XMLStreamException e) {
                throw new MigrationClientException("Error while reading identity.xml", e);
            }
        }
    }

    private static QName getQNameWithIdentityNS(String localPart) {

        return new QName(IdentityCoreConstants.IDENTITY_DEFAULT_NAMESPACE, localPart);
    }
}
