package pl.gda.pg.eti.kask.javaee.enterprise.auth;

import io.jsonwebtoken.*;
import pl.gda.pg.eti.kask.javaee.jsf.business.CryptUtils;
import pl.gda.pg.eti.kask.javaee.jsf.business.boundary.UserService;
import pl.gda.pg.eti.kask.javaee.jsf.business.entities.User;

import javax.enterprise.context.ApplicationScoped;
import javax.enterprise.context.spi.CreationalContext;
import javax.enterprise.inject.spi.Bean;
import javax.enterprise.inject.spi.BeanManager;
import javax.inject.Named;
import javax.naming.InitialContext;
import javax.naming.NamingException;
import javax.security.auth.Subject;
import javax.security.auth.callback.*;
import javax.security.auth.login.CredentialExpiredException;
import javax.security.auth.login.FailedLoginException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;
import java.io.IOException;
import java.security.Principal;
import java.security.acl.Group;
import java.util.List;
import java.util.Map;
import java.util.Set;

@ApplicationScoped
@Named
public class JwtLoginModule implements LoginModule {

    public static final String SIGNING_KEY = "czNjcjN0";

    private Subject subject;
    private CallbackHandler callbackHandler;

    private Principal identity;
    private Group roles;

    private UserService userService;

    @Override
    public void initialize(final Subject subject, final CallbackHandler callbackHandler,
                           final Map<String, ?> sharedState, final Map<String, ?> options) {
        this.subject = subject;
        this.callbackHandler = callbackHandler;

        this.userService = getUserService();
    }

    @Override
    public boolean login() throws LoginException {
        try {

            NameCallback nameCallback = new NameCallback("name:");
            PasswordCallback passwordCallback = new PasswordCallback(" password: ", false);
            callbackHandler.handle(new Callback[]{nameCallback, passwordCallback});
            String password = String.valueOf(passwordCallback.getPassword());

            if(password.isEmpty()){
                return processToken(nameCallback.getName());
            }

            User usr = userService.findUserByCredentials(nameCallback.getName(), CryptUtils.sha256(password));
            if(usr == null)
                return false;

            identity = () -> usr.getLogin();
            roles = new RolesGroup(usr.getRoles());

            return true;

        } catch ( UnsupportedCallbackException |IOException|SignatureException | MalformedJwtException | UnsupportedJwtException | IllegalArgumentException e) {
            throw new FailedLoginException("Invalid token");
        } catch (ExpiredJwtException e) {
            throw new CredentialExpiredException("Token expired");
        }
    }

    private UserService getUserService() {
        try {
            BeanManager beanManager = InitialContext.doLookup( "java:comp/BeanManager" );
            Set<Bean<?>> beans = beanManager.getBeans( UserService.class );
            if ( beans.isEmpty() ) {
                throw new RuntimeException( "Failed looking up CDI Bean " + UserService.class.getName()
                        + ": Found " + beans.size() + " " );
            }
            Bean<?> bean = beans.iterator().next();
            CreationalContext<?> context = beanManager.createCreationalContext( bean );
            return (UserService) beanManager.getReference( bean, UserService.class, context );
        } catch ( NamingException e ) {
            throw new RuntimeException( e );
        }
    }


    private boolean processToken(String jwt){
        Jws<Claims> claims = Jwts.parser().setSigningKey(SIGNING_KEY).parseClaimsJws(jwt);

        identity = getIdentity(claims);
        roles = getRoles(claims);

        return true;
    }

    @Override
    public boolean commit() throws LoginException {
        Set<Principal> principals = subject.getPrincipals();
        principals.add(identity);
        principals.add(roles);
        return true;
    }

    @Override
    public boolean abort() throws LoginException {
        identity = roles = null;
        return true;
    }

    @Override
    public boolean logout() throws LoginException {
        identity = roles = null;
        return true;
    }

    private Principal getIdentity(Jws<Claims> claims) {
        String username = claims.getBody().getSubject();
        return () -> username;
    }

    private Group getRoles(Jws<Claims> claims) {
        List<String> roleNames = claims.getBody().get("roles", List.class);
        RolesGroup roles = new RolesGroup(roleNames);
        return roles;
    }

    private String getJwt() throws LoginException {
        NameCallback nameCallback = new NameCallback("name:");
        PasswordCallback passwordCallback = new PasswordCallback(" password: ", false);


        try {
            callbackHandler.handle(new Callback[]{nameCallback, passwordCallback});
            String name = nameCallback.getName();

            if (name == null) {
                throw new LoginException();
            }

            return name;

        } catch (IOException | UnsupportedCallbackException e) {
            throw new LoginException();
        }
    }
}