package com.oguerreiro.resilient.repository.security;

import java.util.HashMap;
import java.util.Map;

import org.springframework.security.saml2.provider.service.authentication.AbstractSaml2AuthenticationRequest;
import org.springframework.security.saml2.provider.service.web.Saml2AuthenticationRequestRepository;

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

/**
 * <p>
 * This replaces the standard Saml2 in-memory/session request cache. This is actually almost the same, storing the SAML2
 * requests in a HashMap. I was experiencing a problem that SAML2 Request were being lost between SAML request and
 * response (probably because of a Angular error that tries user+password login between calls). Nevertheless, the SAML
 * should work, but didn't.
 * </p>
 * <p>
 * THEN, I tried to implement SAML request persistence in database... WHAT A MESS. SpringSecurity 6.3+ SAML2 really
 * doesn't allow this. The needed AbstractSaml2AuthenticationRequest instance CAN'T be (re)created nor it can be
 * serialized. In the end, I'm unable to implement the repository loadAuthenticationRequest() method and return a valid
 * AbstractSaml2AuthenticationRequest.
 * Quick more stable solution, a custom in-memory/session request cache.
 * </p>
 * <p>
 * But this has limits. If I need a multi-instance server (load-balancing) this will fail. The proposed solution is a
 * distributed cache. It will work almost like this, but instead of a in-memory cache per instance, it will be
 * centralized. When needed, investigate : "centralized distributed cache (Redis, Hazelcast, etc.)"
 * </p>
 * <p>
 * NOTE: This is NOT a bean it self (@Component) because I need to register this in a special way for the SpringSecurity
 * SAML2 auto-binds. In SecurityConfiguration a bean is declared this way:
 * 
 * <pre>
 *   <code>
  &#64;Bean
  Saml2AuthenticationRequestRepository<AbstractSaml2AuthenticationRequest> authenticationRequestRepository() {
    return new ResilientSaml2AuthenticationRequestRepository();
  }
 *   </code>
 * </pre>
 * </p>
 */
//@Component - KEEP commented
public class ResilientSaml2AuthenticationRequestRepository
    implements Saml2AuthenticationRequestRepository<AbstractSaml2AuthenticationRequest> {

  // TODO: Add the creation timestamp, to improve security. Like a 5 second time resolver. Between initiation (request) and authentication (response)
  private Map<String, AbstractSaml2AuthenticationRequest> inMemory = new HashMap<String, AbstractSaml2AuthenticationRequest>();

  public ResilientSaml2AuthenticationRequestRepository() {
  }

  @Override
  public AbstractSaml2AuthenticationRequest loadAuthenticationRequest(HttpServletRequest request) {
    String relayState = request.getParameter("RelayState");
    if (relayState != null) {
      AbstractSaml2AuthenticationRequest authRequest = inMemory.get(relayState);

      return authRequest;
    }

    return null;
  }

  @Override
  public void saveAuthenticationRequest(AbstractSaml2AuthenticationRequest authenticationRequest,
      HttpServletRequest request, HttpServletResponse response) {
    String relayState = authenticationRequest.getRelayState();
    if (relayState != null) {
      inMemory.put(relayState, authenticationRequest);
    }
  }

  @Override
  public AbstractSaml2AuthenticationRequest removeAuthenticationRequest(HttpServletRequest request,
      HttpServletResponse response) {
    String relayState = request.getParameter("RelayState");
    if (relayState != null) {
      inMemory.remove(relayState);
    }
    return null;
  }

}
