Untitled

package dev.lamian.apps.providers;

import dev.lamian.apps.model.UserSchema;
import dev.lamian.apps.utils.JpaUtil;
import jakarta.persistence.EntityManager;
import jakarta.persistence.NoResultException;
import jakarta.persistence.Query;
import lombok.extern.slf4j.Slf4j;
import org.keycloak.component.ComponentModel;
import org.keycloak.connections.jpa.JpaConnectionProvider;
import org.keycloak.credential.CredentialInput;
import org.keycloak.credential.CredentialInputValidator;
import org.keycloak.models.GroupModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.models.UserModel;
import org.keycloak.storage.UserStorageProvider;
import org.keycloak.storage.user.UserLookupProvider;
import org.keycloak.storage.user.UserQueryProvider;


import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

@Slf4j
public class PostgresUserStorageProvider implements UserStorageProvider, UserLookupProvider, CredentialInputValidator, UserQueryProvider {
    private final KeycloakSession session;
    private final ComponentModel model;
    private final EntityManager em;

    public PostgresUserStorageProvider(KeycloakSession session, ComponentModel model) {
        System.getenv().forEach((key, value) ->
                System.out.println("PostgresUserStorageProvider: " + key + " = " + value));
        this.model = model;
        this.session = session;
        this.em = session.getProvider(JpaConnectionProvider.class).getEntityManager();
    }

    @Override
    public boolean supportsCredentialType(String s) {
        return s.equals("password");
    }

    @Override
    public boolean isConfiguredFor(RealmModel realmModel, UserModel userModel, String s) {
        return s.equals("password");
    }

    @Override
    public boolean isValid(RealmModel realmModel, UserModel userModel, CredentialInput credentialInput) {
//        EntityManager em = JpaUtil.getInstance().getEntityManager();
        try {
            UserSchema userEntity = em.createQuery("SELECT u FROM UserSchema u WHERE u.email = :email", UserSchema.class)
                    .setParameter("email", userModel.getUsername())
                    .getSingleResult();

            return credentialInput.getChallengeResponse().equals(userEntity.getPassword());
        } catch (NoResultException e) {
            return false;
        } finally {
            em.close();
        }
    }

    @Override
    public void close() {
        JpaUtil.getInstance().close();
    }

    @Override
    public UserModel getUserById(RealmModel realmModel, String s) {
        return null;
    }

    @Override
    public UserModel getUserByUsername(RealmModel realmModel, String s) {
        return getUserByEmail(realmModel, s);
    }

    @Override
    public UserModel getUserByEmail(RealmModel realmModel, String s) {
//        EntityManager em = JpaUtil.getInstance().getEntityManager();
        try {
            UserSchema userEntity = em.createQuery("SELECT u FROM UserSchema u WHERE u.email = :email",
                            UserSchema.class)
                    .setParameter("email", s)
                    .getSingleResult();
            return createUserModel(realmModel, userEntity);
        } catch (NoResultException e) {
            return null;
        } finally {
            em.close();
        }
    }

    @Override
    public int getUsersCount(RealmModel realm) {
        log.info("getUsersCount {}", realm.getName());
        try {
            Query query = em.createQuery("SELECT COUNT(u) FROM UserSchema u");
            log.info("getUsersCount {}", query.getSingleResult());
            return ((Number) query.getSingleResult()).intValue();
        } catch (Exception e) {
            log.error("getUsersCount error {}", e.getMessage());
            return 0;
        }
    }

    private UserModel createUserModel(RealmModel realmModel, UserSchema u) {
        UserModel userModel = session.users().addUser(realmModel, u.getEmail());
        userModel.setEnabled(true);
        userModel.setEmail(u.getEmail());
        userModel.setFirstName(u.getName());
        userModel.setUsername(u.getEmail());
        log.info("createUserModel, {}", userModel.getEmail());
        return userModel;
    }

    @Override
    public Stream<UserModel> searchForUserStream(RealmModel realm, Map<String, String> params) {
        log.info("Searching for users in realm 1: {}", realm.getName());
        return UserQueryProvider.super.searchForUserStream(realm, params);
    }

    @Override
    public Stream<UserModel> searchForUserStream(RealmModel realm, Map<String, String> params, Integer firstResult, Integer maxResults) {
        log.info("Searching for users in realm 2: {}", params);
        String search = params.get("search");

        // Xây dựng truy vấn JPA
        String jpql = "SELECT u FROM UserSchema u";
        if (search != null && !search.trim().isEmpty()) {
            jpql += " WHERE LOWER(u.email) LIKE :search";
        }

        jakarta.persistence.Query query = em.createQuery(jpql, UserSchema.class);

        if (search != null && !search.trim().isEmpty()) {
            query.setParameter("search", "%" + search.toLowerCase().trim() + "%");
        }

        if (firstResult != null) {
            query.setFirstResult(firstResult);
        }
        if (maxResults != null) {
            query.setMaxResults(maxResults);
        }

        List<UserSchema> users = query.getResultList();
        log.info("Found {} users", users.size());
        return users.stream()
                .map(user -> createUserModel(realm, user));
    }

    @Override
    public Stream<UserModel> getGroupMembersStream(RealmModel realm, GroupModel group, Integer firstResult, Integer maxResults) {
        return Stream.empty();
    }

    @Override
    public Stream<UserModel> searchForUserByUserAttributeStream(RealmModel realm, String attrName, String attrValue) {
        return Stream.empty();
    }
}