package com.airlock.iam.saml2.infrastructure.sp;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.saml2.domain.model.Saml2Binding;
import com.airlock.iam.saml2.domain.model.Saml2RelayState;
import com.airlock.iam.saml2.domain.model.idp.Saml2IdpEntityId;
import com.airlock.iam.saml2.domain.model.idp.Saml2RequestParameter;
import com.airlock.iam.saml2.domain.model.idp.Saml2ResponseParameter;
import com.airlock.iam.saml2.domain.model.sp.Saml2AuthnRequestId;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpEntityId;
import com.airlock.iam.saml2.domain.model.sp.Saml2SpMetaAlias;
import com.airlock.iam.saml2.infrastructure.request.Saml2Request;
import com.airlock.iam.saml2.infrastructure.request.Saml2Response;
import com.airlock.iam.saml2.infrastructure.service.sp.Saml2SpFederationLibrary;
import com.sun.identity.plugin.session.SessionException;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.profile.LogoutRequestDataContainer;
import com.sun.identity.saml2.profile.LogoutUtil;
import com.sun.identity.saml2.profile.SPACSUtils;
import com.sun.identity.saml2.profile.SPCache;
import com.sun.identity.saml2.profile.SPSSOFederate;
import com.sun.identity.saml2.profile.SPSingleLogout;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.net.URI;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

public final class Saml2SpFederationLibraryImpl implements Saml2SpFederationLibrary {
   private final Provider<Saml2Request> request;
   private final Provider<Saml2Response> response;

   @Inject
   Saml2SpFederationLibraryImpl(Provider<Saml2Request> request, Provider<Saml2Response> response) {
      this.request = request;
      this.response = response;
   }

   @Override
   public AssertionData processAssertionFromSamlResponse() {
      try {
         return AssertionDataImpl.fromResponseMap(
            SPACSUtils.processResponseForFedlet((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get())
         );
      } catch (SAML2Exception | IOException | SessionException | ServletException e) {
         throw new FailureException("Error processing assertion from IdP: " + e.getMessage(), e);
      }
   }

   @Override
   public Saml2AuthnRequestId generateNewSpInitiatedSsoRequestId() {
      return Saml2AuthnRequestId.parseAuthnRequestId(SAML2SDKUtils.generateID());
   }

   @Override
   public void initiateAuthnRequest(
      Saml2SpEntityId spEntityId, Saml2IdpEntityId idpEntityId, String realmName, Saml2AuthnRequestId requestId, Saml2Binding binding
   ) {
      try {
         SPSSOFederate.initiateAuthnRequest(
            (HttpServletRequest)this.request.get(),
            (HttpServletResponse)this.response.get(),
            spEntityId.value(),
            idpEntityId.value(),
            realmName,
            Map.of("reqBinding", List.of(binding.value())),
            requestId.value()
         );
      } catch (SAML2Exception e) {
         throw new FailureException(e);
      }
   }

   @Override
   public LogoutRequestDataContainer validateIdpInitiatedSloRequest() {
      Optional<Saml2RelayState> relayState = ((Saml2Request)this.request.get()).relayStateParameter();
      Optional<Saml2RequestParameter> samlRequest = ((Saml2Request)this.request.get()).samlRequest();
      if (samlRequest.isEmpty()) {
         throw new ViolationException("No SAMLRequest parameter found in IdP-Initiated SLO request");
      }

      try {
         return SPSingleLogout.validateLogoutRequest(
               (HttpServletRequest)this.request.get(), samlRequest.get().value(), relayState.map(Saml2RelayState::value).orElse(null)
            )
            .orElseThrow(() -> new ViolationException("Invalid SAMLRequest"));
      } catch (SAML2Exception | SessionException e) {
         throw new FailureException("Error while processing Single-Logout Request from IdP: " + e.getMessage(), e);
      }
   }

   @Override
   public void continueIdpInitiatedSlo(LogoutRequestDataContainer logoutData) {
      try {
         SPSingleLogout.handleLogoutRequest((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get(), logoutData);
      } catch (SAML2Exception e) {
         throw new FailureException("IdP-initiated SLO: Error while processing Single-Logout continuation request " + e.getMessage(), e);
      }
   }

   @Override
   public boolean initiateLogoutRequest(Saml2SpMetaAlias metaAlias, Saml2IdpEntityId idpEntityId, URI relayStateUri) {
      try {
         String binding = LogoutUtil.getSLOBindingInfo((HttpServletRequest)this.request.get(), metaAlias.value(), "SPRole", idpEntityId.value());
         Map<String, String> paramsMap = new HashMap<>();
         paramsMap.put("spMetaAlias", metaAlias.value());
         paramsMap.put("role", "SPRole");
         paramsMap.put("binding", binding);
         paramsMap.put("RelayState", relayStateUri.toASCIIString());
         SPSingleLogout.initiateLogoutRequest((HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get(), binding, paramsMap);
         return ((Saml2Response)this.response.get()).isHandledBySaml2Library();
      } catch (SAML2Exception e) {
         throw new FailureException("Error while initiating SLO request in library: " + e.getMessage(), e);
      }
   }

   @Override
   public Optional<URI> relayStateUriFromLogoutResponse() {
      Optional<Saml2RelayState> rawRelayState = ((Saml2Request)this.request.get()).relayStateParameter();
      Optional<Saml2RelayState> relayStateFromCache = rawRelayState.map(Saml2RelayState::toString)
         .map(SPCache.relayStateHash()::remove)
         .map(Saml2RelayState::relayState);
      return relayStateFromCache.isPresent() ? relayStateFromCache.flatMap(Saml2RelayState::toUri) : rawRelayState.flatMap(Saml2RelayState::toUri);
   }

   @Override
   public void processSpInitiatedLogoutResponse(URI relayStateUri) {
      Optional<Saml2ResponseParameter> samlResponse = ((Saml2Request)this.request.get()).samlResponse();
      if (samlResponse.isEmpty()) {
         throw new ViolationException("No SAMLResponse parameter found in SP-Initiated SLO responder request");
      }

      try {
         SPSingleLogout.processLogoutResponse(
            (HttpServletRequest)this.request.get(), (HttpServletResponse)this.response.get(), samlResponse.get().value(), relayStateUri.toASCIIString()
         );
      } catch (SAML2Exception | SessionException e) {
         throw new FailureException("Error while processing SLO response in library: " + e.getMessage(), e);
      }
   }
}