package cz.muni.ics.oidc.saml;

import cz.muni.ics.oauth2.model.ClientDetailsEntity;
import cz.muni.ics.oauth2.model.DeviceCode;
import cz.muni.ics.oauth2.repository.impl.DeviceCodeRepository;
import cz.muni.ics.oauth2.service.ClientDetailsEntityService;
import cz.muni.ics.oidc.models.Facility;
import cz.muni.ics.oidc.models.PerunAttributeValue;
import cz.muni.ics.oidc.server.adapters.PerunAdapter;
import cz.muni.ics.oidc.server.configurations.FacilityAttrsConfig;
import cz.muni.ics.oidc.server.configurations.PerunOidcConfig;
import cz.muni.ics.oidc.server.filters.AuthProcFilterConstants;
import lombok.extern.slf4j.Slf4j;
import org.opensaml.common.SAMLException;
import org.opensaml.saml2.metadata.AssertionConsumerService;
import org.opensaml.saml2.metadata.SPSSODescriptor;
import org.opensaml.saml2.metadata.provider.MetadataProviderException;
import org.opensaml.ws.message.encoder.MessageEncodingException;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.saml.SAMLConstants;
import org.springframework.security.saml.SAMLEntryPoint;
import org.springframework.security.saml.context.SAMLMessageContext;
import org.springframework.security.saml.util.SAMLUtil;
import org.springframework.security.saml.websso.WebSSOProfileOptions;
import org.springframework.util.StringUtils;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static cz.muni.ics.oauth2.web.endpoint.DeviceEndpoint.PATH_DEVICE_AUTHORIZE;
import static cz.muni.ics.oauth2.web.endpoint.DeviceEndpoint.USER_CODE;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.AARC_IDP_HINT;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.BLOCKED_IDPS_ACR_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.CLIENT_ID_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.EFILTER_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.FILTER_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.IDP_ENTITY_ID_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.ONLY_ALLOWED_IDPS_ACR_PREFIX;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.PARAM_CLIENT_ID;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.PARAM_MAX_AGE;
import static cz.muni.ics.oidc.server.filters.AuthProcFilterConstants.PARAM_PROMPT;

@Slf4j
public class PerunSamlEntryPoint extends SAMLEntryPoint {

    private final PerunAdapter perunAdapter;
    private final PerunOidcConfig config;
    private final FacilityAttrsConfig facilityAttrsConfig;
    private final SamlProperties samlProperties;
    private final DeviceCodeRepository deviceCodeRepository;
    private final ClientDetailsEntityService clientDetailsEntityService;

    public PerunSamlEntryPoint(PerunAdapter perunAdapter,
                               PerunOidcConfig config,
                               FacilityAttrsConfig facilityAttrsConfig,
                               SamlProperties samlProperties,
                               DeviceCodeRepository deviceCodeRepository,
                               ClientDetailsEntityService clientDetailsEntityService
    ) {
        this.perunAdapter = perunAdapter;
        this.config = config;
        this.facilityAttrsConfig = facilityAttrsConfig;
        this.samlProperties = samlProperties;
        this.deviceCodeRepository = deviceCodeRepository;
        this.clientDetailsEntityService = clientDetailsEntityService;
    }

    @Override
    public void commence(HttpServletRequest request, HttpServletResponse response, AuthenticationException e)
            throws IOException, ServletException {
        try {
            SAMLMessageContext context = contextProvider.getLocalAndPeerEntity(request, response);

            Map<String, String> parameters = new HashMap<>();
            Iterator it = request.getParameterNames().asIterator();
            while(it.hasNext()) {
                String param = (String) it.next();
                parameters.put(param, request.getParameter(param));
            }
            processDeviceCodeRequestParameters(request, parameters);
            if (isECP(context)) {
                initializeECP(parameters, context, e);
            } else if (isDiscovery(context)) {
                initializeDiscovery(context);
            } else {
                String requestUrl = request.getRequestURL().toString() + '?' + request.getQueryString();
                initializeSSO(parameters, context, e, requestUrl);
            }
        } catch (SAMLException | MetadataProviderException | MessageEncodingException e1) {
            log.debug("Error initializing entry point", e1);
            throw new ServletException(e1);
        }
    }

    protected WebSSOProfileOptions getExtendedOptions(Map<String, String> requestParameters,
                                                      SAMLMessageContext context,
                                                      AuthenticationException exception)
            throws MetadataProviderException
    {
        WebSSOProfileOptions options = super.getProfileOptions(context, exception);
        addExtraParams(requestParameters, options);
        return options;
    }


    // copy from super class to call our getProfileOptions
    protected void initializeECP(Map<String, String> requestParameters, SAMLMessageContext context,
                                 AuthenticationException e)
            throws MetadataProviderException, SAMLException, MessageEncodingException
    {
        WebSSOProfileOptions options = getExtendedOptions(requestParameters, context, e);
        log.debug("Processing SSO using ECP profile");
        webSSOprofileECP.sendAuthenticationRequest(context, options);
        samlLogger.log(SAMLConstants.AUTH_N_REQUEST, SAMLConstants.SUCCESS, context);
    }

    // copy from super class to call our getProfileOptions
    protected void initializeSSO(Map<String, String> requestParameters, SAMLMessageContext context,
                                 AuthenticationException e, String originalUrl)
            throws MetadataProviderException, SAMLException, MessageEncodingException
    {
        WebSSOProfileOptions options = getExtendedOptions(requestParameters, context, e);
        AssertionConsumerService consumerService = SAMLUtil.getConsumerService(
                (SPSSODescriptor) context.getLocalEntityRoleMetadata(), options.getAssertionConsumerIndex());

        // HoK WebSSO
        if (SAMLConstants.SAML2_HOK_WEBSSO_PROFILE_URI.equals(consumerService.getBinding())) {
            if (webSSOprofileHoK == null) {
                log.warn("WebSSO HoK profile was specified to be used, but profile is not " +
                    "configured in the EntryPoint, HoK will be skipped");
            } else {
                log.debug("Processing SSO using WebSSO HolderOfKey profile");
                webSSOprofileHoK.sendAuthenticationRequest(context, options);
                samlLogger.log(SAMLConstants.AUTH_N_REQUEST, SAMLConstants.SUCCESS, context);
                return;
            }
        }

        context = processAarcIdpHint(context, requestParameters.getOrDefault(AARC_IDP_HINT, null));
        context.setRelayState(originalUrl);
        // Ordinary WebSSO
        log.debug("Processing SSO using WebSSO profile");
        webSSOprofile.sendAuthenticationRequest(context, options);
        samlLogger.log(SAMLConstants.AUTH_N_REQUEST, SAMLConstants.SUCCESS, context);
    }

    private void processDeviceCodeRequestParameters(HttpServletRequest request,
                                                                   Map<String, String> parameters)
    {
        if (PATH_DEVICE_AUTHORIZE.equals(request.getServletPath())) {
            DeviceCode dc = deviceCodeRepository.getByUserCode(request.getParameter(USER_CODE));
            parameters.putAll(dc.getRequestParameters());
        }
    }

    private SAMLMessageContext processAarcIdpHint(SAMLMessageContext context, String aarcIdpHint) {
        if (StringUtils.hasText(aarcIdpHint)) {
            PerunSAMLMessageContext messageContext = new PerunSAMLMessageContext(context);
            messageContext.setAarcIdpHint(aarcIdpHint);
            context = messageContext;
            log.debug("Added aarc_idp_hint to context '{}'", context);
        }
        return context;
    }

    private void addExtraParams(Map<String, String> requestParameters, WebSSOProfileOptions options) {
        log.debug("Transforming OIDC params to SAML");
        processPrompt(requestParameters, options);
        processAcrValues(requestParameters, options);
        addRequesterIds(requestParameters, options);
        processMaxAge(requestParameters, options);
    }

    private void processPrompt(Map<String, String> requestParameters, WebSSOProfileOptions options) {
        String prompt = requestParameters.getOrDefault(PARAM_PROMPT, "");
        if (PerunSamlUtils.needsReAuthByPrompt(prompt)) {
            log.debug("Transformed prompt parameter ({}) to SAML forceAuthn=true", prompt);
            options.setForceAuthN(true);
        }
        if ("none".equalsIgnoreCase(prompt)) {
            log.debug("Detected prompt=none, translating to 'isPassive=true' in SAML");
            options.setPassive(true);
        }
    }

    private void processMaxAge(Map<String, String> requestParameters, WebSSOProfileOptions options) {
        String maxAgeStr = requestParameters.getOrDefault(PARAM_MAX_AGE, null);
        if (StringUtils.hasText(maxAgeStr) && "0".equals(maxAgeStr)) {
            log.debug("Detected max_age=0, translating to SAML forceAuthn=true");
            options.setForceAuthN(true);
        }
    }

    private void processAcrValues(Map<String, String> requestParameters, WebSSOProfileOptions options) {
        String acrValues = requestParameters.getOrDefault(AuthProcFilterConstants.PARAM_ACR_VALUES, null);
        List<String> acrs = new ArrayList<>();
        if (acrValues != null) {
            log.debug("Processing acr_values parameter: {}", acrValues);
            acrs = convertAcrValuesToList(acrValues);
        }

        String clientId = requestParameters.getOrDefault(AuthProcFilterConstants.PARAM_CLIENT_ID, null);
        if (StringUtils.hasText(clientId)) {
            // ADD FILTER AND E-FILTER
            if (config.isAskPerunForIdpFiltersEnabled() && !hasAcrForcingIdp(acrs)) {
                String idpFilter = extractIdpFilterForRp(clientId);
                if (idpFilter != null) {
                    log.debug("Added IdP filter as SAML AuthnContextClassRef ({})", idpFilter);
                    acrs.add(idpFilter);
                }
            }

            ClientDetailsEntity client = clientDetailsEntityService.loadClientByClientId(clientId);
            if (client != null) {
                // ADD BLOCKED IdPs
                String blockedIdps = getBlockedIdpsAcr(client);
                log.debug("blockedIdps ({})", blockedIdps);
                if (StringUtils.hasText(blockedIdps)) {
                    String acr = BLOCKED_IDPS_ACR_PREFIX + blockedIdps;
                    log.debug("Added blockedIdps as SAML AuthnContextClassRef ({})", acr);
                    acrs.add(acr);
                }

                // ADD ONLY ALLOWED IdPs
                String onlyAllowedIdps = getOnlyAllowedIdpsAcr(client);
                log.debug("allowedIdps ({})", onlyAllowedIdps);
                if (StringUtils.hasText(onlyAllowedIdps)) {
                    String acr = ONLY_ALLOWED_IDPS_ACR_PREFIX + onlyAllowedIdps;
                    log.debug("Added onlyAllowedIdps as SAML AuthnContextClassRef ({})", acr);
                    acrs.add(acr);
                }
            }

            // ADD CLIENT_ID
            if (config.isAddClientIdToAcrs()) {
                String clientIdAcr = CLIENT_ID_PREFIX + requestParameters.get(PARAM_CLIENT_ID);
                log.debug("Adding client_id ACR ({}) to list of AuthnContextClassRefs for purposes" +
                        " of displaying service name on the wayf", clientIdAcr);
                acrs.add(clientIdAcr);
            }
        }

        if (!acrs.isEmpty()) {
            processAcrs(acrs);
            options.setAuthnContexts(acrs);
            log.debug("Transformed acr_values ({}) to SAML AuthnContextClassRef ({})",
                    acrValues, options.getAuthnContexts());
        }
    }

    private void addRequesterIds(Map<String, String> requestParameters, WebSSOProfileOptions options) {
        log.debug("Request parameters: {}", requestParameters);
        String prefix = config.getRequesterIdPrefix();
        if (!StringUtils.hasText(prefix)) {
            return;
        }
        String clientId = requestParameters.getOrDefault(PARAM_CLIENT_ID, null);
        if (StringUtils.hasText(clientId)) {
            log.debug("Adding ClientID ({}) to SAML RequesterIDs", clientId);
            Set<String> requesterIds = options.getRequesterIds();
            if (requesterIds == null) {
                requesterIds = new HashSet<>();
            }
            requesterIds.add(prefix + ':' + clientId);
            options.setRequesterIds(requesterIds);
        } else {
            log.debug("Did not add client_id to SAML RequesterIDs, request parameter client_id not found" +
                    " or has an empty value");
        }
    }

    private void processAcrs(List<String> acrs) {
        if (acrs == null || acrs.isEmpty()) {
            return;
        }

        String[] reservedAcrsPrefixes = samlProperties.getAcrReservedPrefixes();
        Set<String> reservedPrefixes = (reservedAcrsPrefixes != null) ?
                new HashSet<>(Arrays.asList(reservedAcrsPrefixes)) : new HashSet<>();
        if (reservedPrefixes.isEmpty()) {
            return;
        }

        boolean hasNonReserved = false;
        for (String prefix: reservedPrefixes) {
            for (String acr: acrs) {
                if (!acr.startsWith(prefix)) {
                    log.debug("ACR with non reserved prefix found: {}", acr);
                    hasNonReserved = true;
                    break;
                }
            }
            if (hasNonReserved) {
                break;
            }
        }

        if (!hasNonReserved) {
            List<String> toBeAdded = new LinkedList<>(Arrays.asList(samlProperties.getAcrsToBeAdded()));
            log.debug("NO ACR with non reserved prefix found, adding following: {}", toBeAdded);
            acrs.addAll(toBeAdded);
        }
    }

    private List<String> convertAcrValuesToList(String acrValues) {
        List<String> acrs = new LinkedList<>();
        if (StringUtils.hasText(acrValues)) {
            String[] parts = acrValues.split(" ");
            for (String acr : parts) {
                if (StringUtils.hasText(acr)) {
                    acrs.add(acr);
                }
            }
        }
        return acrs;
    }

    private boolean hasAcrForcingIdp(List<String> acrs) {
        boolean hasIdpEntityId = acrs != null
            && !acrs.isEmpty()
            && acrs.stream().anyMatch(
            acr -> StringUtils.hasText(acr) && acr.startsWith(IDP_ENTITY_ID_PREFIX)
        );

        if (hasIdpEntityId) {
            log.debug("Request contains ACR to force specific IdP, no configured IdP filter will be used");
        }
        return hasIdpEntityId;
    }

    private String extractIdpFilterForRp(String clientId) {
        if (!config.isAskPerunForIdpFiltersEnabled()) {
            return null;
        }
        Facility facility = null;
        if (clientId != null) {
            facility = perunAdapter.getFacilityByClientId(clientId);
        }
        Map<String, PerunAttributeValue> filterAttributes = new HashMap<>();
        if (facility != null) {
            filterAttributes = this.getFacilityFilterAttributes(facility);
        }

        String idpFilterAcr = null;
        String idpFilter = fetchRpIdpFilter(filterAttributes, facilityAttrsConfig.getWayfFilterAttr());
        if (StringUtils.hasText(idpFilter)) {
            idpFilterAcr = FILTER_PREFIX + idpFilter;
        } else {
            String idpEFilter = fetchRpIdpFilter(filterAttributes, facilityAttrsConfig.getWayfEFilterAttr());
            if (StringUtils.hasText(idpEFilter)) {
                idpFilterAcr = EFILTER_PREFIX + idpEFilter;
            }
        }
        return idpFilterAcr;
    }

    private Map<String, PerunAttributeValue> getFacilityFilterAttributes(Facility facility) {
        if (facility != null && facility.getId() != null) {
            List<String> attrsToFetch = new ArrayList<>();
            attrsToFetch.add(facilityAttrsConfig.getWayfEFilterAttr());
            attrsToFetch.add(facilityAttrsConfig.getWayfFilterAttr());
            return perunAdapter.getFacilityAttributeValues(facility, attrsToFetch);
        }
        return new HashMap<>();
    }

    private String fetchRpIdpFilter(Map<String, PerunAttributeValue> filterAttributes, String attrName) {
        String result = null;
        if (filterAttributes.get(attrName) != null) {
            PerunAttributeValue filterAttribute = filterAttributes.get(attrName);
            if (filterAttribute != null && filterAttribute.valueAsString() != null) {
                result = filterAttribute.valueAsString();
            }
        }
        return result;
    }

    private String getOnlyAllowedIdpsAcr(ClientDetailsEntity client) {
        String result = null;
        if (config.isOnlyAllowedIdpsEnabled()) {
            Set<String> idps = client.getOnlyAllowedIdps();
            if (idps != null && !idps.isEmpty()) {
                result = String.join(";", idps);
            }
        }
        return result;
    }

    private String getBlockedIdpsAcr(ClientDetailsEntity client) {
        String result = null;
        if (config.isBlockedIdpsEnabled()) {
            Set<String> idps = client.getBlockedIdps();
            if (idps != null && !idps.isEmpty()) {
                result = String.join(";", idps);
            }
        }
        return result;
    }

}
