/**
 * Copyright (c) 2024, WSO2 LLC. (https://www.wso2.com).
 *
 * 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.dao;

import org.wso2.carbon.identity.application.common.model.IdPGroup;
import org.wso2.carbon.identity.core.migrate.MigrationClientException;
import org.wso2.carbon.identity.core.util.IdentityDatabaseUtil;
import org.wso2.carbon.is.migration.service.v700.model.IdPRole;
import org.wso2.carbon.user.core.UserCoreConstants;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

public class IdPDAO {

    private static final String GET_IDP_ROLES_BY_TENANT_ID = "SELECT ID, IDP_ID, ROLE FROM IDP_ROLE WHERE " +
            "TENANT_ID = ?";
    private static final String ADD_IDP_GROUPS_SQL = "INSERT INTO IDP_GROUP (IDP_ID, TENANT_ID, GROUP_NAME, UUID) "
            + "VALUES (?, ?, ?, ?)";
    private static final String GET_IDP_GROUP = "SELECT ID, IDP_ID, TENANT_ID, GROUP_NAME, UUID FROM IDP_GROUP WHERE "
            + "TENANT_ID = ? AND IDP_ID = ? AND GROUP_NAME = ?";
    private static final String GET_IDP_ROLE_MAPPINGS_BY_IDP_ROLE_ID = "SELECT USER_STORE_ID, LOCAL_ROLE FROM " +
            "IDP_ROLE_MAPPING WHERE TENANT_ID = ? AND IDP_ROLE_ID = ?";
    private static final String ADD_IDP_GROUP_TO_ROLE_ASSIGNMENT = "INSERT INTO UM_IDP_GROUP_ROLE (UM_ROLE_ID, " +
            "UM_GROUP_ID, UM_TENANT_ID) VALUES (?, ?, ?)";

    public List<IdPRole> getIdPRoles(int tenantId) throws MigrationClientException {

        List<IdPRole> idPRoles = new ArrayList<>();
        try (Connection connection = IdentityDatabaseUtil.getDBConnection(false);
             PreparedStatement prepStmt = connection.prepareStatement(GET_IDP_ROLES_BY_TENANT_ID)) {
            prepStmt.setInt(1, tenantId);
            ResultSet resultSet = prepStmt.executeQuery();
            while (resultSet.next()) {
                IdPRole idPRole = new IdPRole();
                idPRole.setId(resultSet.getInt("ID"));
                idPRole.setIdPId(resultSet.getInt("IDP_ID"));
                idPRole.setRole(resultSet.getString("ROLE"));
                idPRoles.add(idPRole);
            }
        } catch (SQLException e) {
            throw new MigrationClientException("Error while retrieving IdP roles", e);
        }
        return idPRoles;
    }

    public IdPGroup addIdPGroup(int tenantId, int idPId, String groupName) throws MigrationClientException {

        try (Connection connection = IdentityDatabaseUtil.getDBConnection(true);
             PreparedStatement prepStmt = connection.prepareStatement(ADD_IDP_GROUPS_SQL)) {
            prepStmt.setInt(1, idPId);
            prepStmt.setInt(2, tenantId);
            prepStmt.setString(3, groupName);
            prepStmt.setString(4, UUID.randomUUID().toString());
            prepStmt.execute();
        } catch (SQLException e) {
            throw new MigrationClientException("Error while adding IdP groups", e);
        }
        return getIdPGroup(tenantId, idPId, groupName);
    }

    public String getLocalRoleMappedToIdPRole(int tenantId, int idPRoleId) throws MigrationClientException {

        try (Connection connection = IdentityDatabaseUtil.getDBConnection(false);
             PreparedStatement prepStmt = connection.prepareStatement(GET_IDP_ROLE_MAPPINGS_BY_IDP_ROLE_ID)) {
            prepStmt.setInt(1, tenantId);
            prepStmt.setInt(2, idPRoleId);
            ResultSet resultSet = prepStmt.executeQuery();
            if (resultSet.next()) {
                String userStoreId = resultSet.getString("USER_STORE_ID");
                String roleName = resultSet.getString("LOCAL_ROLE");
                return userStoreId + UserCoreConstants.DOMAIN_SEPARATOR + roleName;
            }
        } catch (SQLException e) {
            throw new MigrationClientException("Error while retrieving IdP role mappings", e);
        }
        return null;
    }

    public void addIdPGroupToRoleAssignment(Connection connection, int roleId, String idPGroupUUID, int tenantId)
            throws MigrationClientException {

        try {
            PreparedStatement prepStmt = connection.prepareStatement(ADD_IDP_GROUP_TO_ROLE_ASSIGNMENT);
            prepStmt.setInt(1, roleId);
            prepStmt.setString(2, idPGroupUUID);
            prepStmt.setInt(3, tenantId);
            prepStmt.execute();
        } catch (SQLException e) {
            throw new MigrationClientException("Error while adding IdP groups", e);
        }
    }

    public IdPGroup getIdPGroup(int tenantId, int idPId, String groupName) throws MigrationClientException {

        try (Connection connection = IdentityDatabaseUtil.getDBConnection(false);
             PreparedStatement prepStmt = connection.prepareStatement(GET_IDP_GROUP)) {
            prepStmt.setInt(1, tenantId);
            prepStmt.setInt(2, idPId);
            prepStmt.setString(3, groupName);
            ResultSet resultSet = prepStmt.executeQuery();
            if (resultSet.next()) {
                IdPGroup idPGroup = new IdPGroup();
                idPGroup.setIdpGroupName(resultSet.getString("GROUP_NAME"));
                idPGroup.setIdpGroupId(resultSet.getString("UUID"));
                return idPGroup;
            }
        } catch (SQLException e) {
            throw new MigrationClientException("Error while retrieving IdP group", e);
        }
        return null;
    }
}
