package com.oguerreiro.resilient.security.saml2;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.env.Environment;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.saml2.provider.service.authentication.Saml2AuthenticatedPrincipal;
import org.springframework.security.saml2.provider.service.authentication.Saml2Authentication;
import org.springframework.security.web.authentication.AuthenticationFailureHandler;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.stereotype.Component;

import com.oguerreiro.resilient.domain.Organization;
import com.oguerreiro.resilient.domain.User;
import com.oguerreiro.resilient.repository.OrganizationRepository;
import com.oguerreiro.resilient.repository.UserRepository;
import com.oguerreiro.resilient.repository.security.SecurityGroupRepository;
import com.oguerreiro.resilient.security.custom.ResilientUserDetails;
import com.oguerreiro.resilient.security.resillient.SecurityGroup;
import com.oguerreiro.resilient.security.saml2.ResilientSaml2Properties.Registration;

import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;

@Component
public class Saml2AuthenticationHandler implements AuthenticationSuccessHandler, AuthenticationFailureHandler {
  protected final Logger log = LoggerFactory.getLogger(this.getClass());

  @Autowired
  private Environment environment;

  @Autowired
  private OrganizationRepository organizationRepository;

  @Autowired
  private SecurityGroupRepository securityGroupRepository;

  @Autowired
  private ResilientSaml2Properties resilientSaml2Properties;

  @Autowired
  private UserRepository userRepository;

  /**
   * AuthenticationSuccessHandler implementation
   */
  @Override
  public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response,
      Authentication authentication) throws IOException, ServletException {
    String samlResponse = request.getParameter("SAMLResponse");
    String samlXMLResponse = Saml2ResponseLoggingFilter.decodeSamlResponse(samlResponse);

    // Declarations
    Collection<? extends GrantedAuthority> authorities;
    Organization userOrganization = null;
    SecurityGroup securityGroup = null;

    // read SAML2 attributes
    Saml2AuthenticatedPrincipal principal = (Saml2AuthenticatedPrincipal) authentication.getPrincipal();
    Saml2AuthenticationAttributes attr = this.readSaml2Attributes(principal);

    // Search user. If the user (by username) is registered in Resilient, load definitions from Resilient and NOT the ones received from SAML2
    String username = attr.getUsername();
    if (username == null | username.isBlank()) {
      // This is NOT allowed. A username must be provided. If from SAML2, most probably, will be an email
      log.error("The attribute username is expected for SAML2 authentication, but it was NULL or Empty.");
      this.invalidateLogin(request, response);
      return;
    }

    User user = userRepository.findOneWithAuthoritiesByLogin(username).orElse(null);
    if (user != null) {
      // This user can login with SAML2 ?
      if (!user.getAllowSamlAuthentication()) {
        log.error("The user '" + username + "' it's not allowed to login with SAML2 authentication.");
        this.invalidateLogin(request, response);
        return;
      }

      //@formatter:off
      // User found. Setup permissions from resilient config
      authorities = user.getAuthorities().stream()
                      .map(auth -> new SimpleGrantedAuthority(auth.getName()))
                      .collect(Collectors.toList());
      //@formatter:on
      securityGroup = user.getSecurityGroup();
      userOrganization = organizationRepository.findOneByUser(user).orElse(null);
      if (userOrganization == null) {
        // This is NOT allowed. Bad user configuration.
        log.error(
            "The username is registered in Resilient, but without reference to an Organization. This is a configuration error.");
        this.invalidateLogin(request, response);
        return;
      }
    } else {
      // User not found? Use SAML2 values or defaults.
      if (attr.getRole() == null || attr.getRole().isBlank()) {
        // This is NOT allowed. SAML2 or Defaults must provide a Role.
        log.error(
            "A ROLE must be given to a user, but none was calculated. Nor by SAML2 XML, nor defined in defaults.");
        this.invalidateLogin(request, response);
        return;
      }
      authorities = List.of(attr.getRole()).stream().map(SimpleGrantedAuthority::new).collect(Collectors.toList());

      // Organization not found ?
      if (attr.getOrganizationCode() == null || attr.getOrganizationCode().isBlank()) {
        // This is NOT allowed. SAML2 or Defaults must provide a Organization code.
        log.error(
            "An OrganizationCode must be given to a user, but none was calculated. Nor by SAML2 XML, nor defined in defaults.");
        this.invalidateLogin(request, response);
        return;
      } else {
        // Find or NULL
        userOrganization = organizationRepository.findOneByCode(attr.getOrganizationCode()).orElse(null);
      }
      if (userOrganization == null) {
        log.error("The Organization code [" + attr.getOrganizationCode() + "] wasn't found in database.");
        this.invalidateLogin(request, response);
        return;
      }

      // SecurityGroup exists?
      if (attr.getSecurityGroupCode() == null || attr.getSecurityGroupCode().isBlank()) {
        // This is NOT allowed. SAML2 or Defaults must provide a SecurityGroup code.
        log.error(
            "A SecurityGroup code must be given to a user, but none was calculated. Nor by SAML2 XML, nor defined in defaults.");
        this.invalidateLogin(request, response);
        return;
      }

      // Find or fail
      securityGroup = securityGroupRepository.findOneByCodeWithEagerRelationships(attr.getSecurityGroupCode()).orElse(
          null);

      if (securityGroup == null) {
        log.error("The SecurityGroup code [" + attr.getSecurityGroupCode() + "] wasn't found in database.");
        this.invalidateLogin(request, response);
        return;
      }
    }

    // Create a ResilientUserDetails and replace Principal
    ResilientUserDetails userdetails = new ResilientUserDetails(username, "MOCK-PWD", authorities, securityGroup,
        userOrganization, "pt-PT");

    Saml2Authentication newAuthentication = new Saml2Authentication(userdetails, samlXMLResponse,
        userdetails.getAuthorities());
    SecurityContextHolder.getContext().setAuthentication(newAuthentication);

    // Send to success URL, if configured
    String successUrl = resilientSaml2Properties.getSuccessUrl();
    if (successUrl != null && !successUrl.isBlank()) {
      // This is mandatory in DEV environment. Optional in PROD because the app url is the same.
      // Even so, I think its a good practice to define the success url
      response.sendRedirect(successUrl);
    }
  }

  private void invalidateLogin(HttpServletRequest request, HttpServletResponse response) throws IOException {
    // Clear context and invalidate session
    SecurityContextHolder.clearContext();
    request.getSession(false).invalidate();

    // Send error response or redirect
    response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Missing or invalid SAML attributes");
  }

  private boolean isDevProfileActive() {
    return Arrays.asList(environment.getActiveProfiles()).contains("dev");
  }

  private String getPrincipalAttribute(Saml2AuthenticatedPrincipal principal, String attribute) {
    List<?> values = principal.getAttribute(attribute);

    if (values == null || values.isEmpty()) {
      return null;
    }

    return principal.getAttribute(attribute).get(0).toString();
  }

  /**
   * AuthenticationFailureHandler implementation
   */
  @Override
  public void onAuthenticationFailure(HttpServletRequest request, HttpServletResponse response,
      AuthenticationException exception) throws IOException, ServletException {

    // This is a sugar-code when in development environment. 
    if (isDevProfileActive()) {
      // If this is a mock-idp, it can provide the parameter 'SAMLDevEnvironmentUrl' 
      //  that gives the base URL to use. This is because in DEV mode usually the 
      //  Angular side runs in localhost:42000 but server-side is in localhost:8080.
      //  Without this, SAMLv2 authentication would end up in error redirecting the user to 
      //  localhost:8080 (NOT the client-side)
      // In PROD we don't need this, because the app url is the same
      String failureUrl = resilientSaml2Properties.getFailureUrl();

      if (failureUrl != null && !failureUrl.isBlank()) {
        response.sendRedirect(failureUrl);
      }
    }
  }

  /**
   * Extract the expected user attributes from SAML2Response. There's some standard expected attributes, those will be
   * mapped to properties, the remaining will be put into a MAP.
   */
  private Saml2AuthenticationAttributes readSaml2Attributes(Saml2AuthenticatedPrincipal principal) {
    // Get the RelyingParty ID (IDP)
    String registrationId = principal.getRelyingPartyRegistrationId();

    // Get registration, from config properties for the RelyingParty ID
    Registration registration = this.resilientSaml2Properties.getRelyingparty().getRegistration().get(registrationId);

    if (registration == null) {
      // This is VERY strange
    }

    Saml2AuthenticationAttributes authAttr = new Saml2AuthenticationAttributes();
    Map<String, String> atributes = registration.getAssertingparty().getAttributes();
    Map<String, String> defaults = registration.getAssertingparty().getDefaults();

    for (Entry<String, String> entry : atributes.entrySet()) {
      String attrValue = this.getPrincipalAttribute(principal, entry.getValue());
      authAttr.addAttribute(entry.getKey(), attrValue, defaults.get(entry.getKey()));
    }

    return authAttr;
  }
}
