package com.airlock.iam.saml2.infrastructure.service.idp;

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.domain.model.logging.Saml2LogTypes;
import com.airlock.iam.saml2.infrastructure.util.Saml2RequestThreadContext;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2SDKUtils;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.protocol.AuthnRequest;
import com.sun.identity.saml2.protocol.ProtocolFactory;
import com.sun.identity.shared.encode.Base64;
import com.sun.identity.shared.xml.XMLUtils;
import jakarta.annotation.Nonnull;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import jakarta.servlet.http.HttpServletRequest;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.Optional;
import org.w3c.dom.Document;

public final class Saml2AuthnRequestExtractorImpl implements Saml2AuthnRequestExtractor {
   private static final ELogger LOG = ELoggerFactory.getLogger(Saml2AuthnRequestExtractorImpl.class);
   private final Provider<HttpServletRequest> request = Saml2RequestThreadContext::getRequest;

   @Override
   public Optional<AuthnRequest> extractAuthnRequestFromCurrentRequest() {
      String samlRequest = ((HttpServletRequest)this.request.get()).getParameter("SAMLRequest");
      if (samlRequest == null) {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "No SAMLRequest parameter in request");
         return Optional.empty();
      }

      if (((HttpServletRequest)this.request.get()).getMethod().equals("GET")) {
         return this.extractFromGetRequest(samlRequest);
      }

      if (((HttpServletRequest)this.request.get()).getMethod().equals("POST")) {
         return this.extractFromPostRequest(samlRequest);
      }

      LOG.warn(
         Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Cannot extract request from " + LogUtil.logSafeQuote(((HttpServletRequest)this.request.get()).getMethod())
      );
      return Optional.empty();
   }

   @Nonnull
   private Optional<AuthnRequest> extractFromGetRequest(String samlRequest) {
      LOG.debug("Receiving SAML2 Request via GET");
      String outputString = SAML2Utils.decodeFromRedirect(samlRequest);
      if (outputString != null) {
         try {
            return Optional.of(ProtocolFactory.getInstance().createAuthnRequest(outputString));
         } catch (SAML2Exception e) {
            LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error parsing AuthnRequest: " + e.getMessage(), e);
         }
      } else {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error decoding AuthnRequest");
      }

      LOG.debug("Invalid AuthnRequest: " + LogUtil.logSafeString(samlRequest));
      return Optional.empty();
   }

   private Optional<AuthnRequest> extractFromPostRequest(String samlRequest) {
      LOG.debug("Receiving SAML2 Request via POST");

      try {
         byte[] raw = Base64.decode(samlRequest);
         if (raw != null) {
            try (ByteArrayInputStream bis = new ByteArrayInputStream(raw)) {
               Document doc = XMLUtils.toDOMDocument(bis, SAML2SDKUtils.DEBUG);
               if (doc != null) {
                  return Optional.of(ProtocolFactory.getInstance().createAuthnRequest(doc.getDocumentElement()));
               }
            } catch (IOException e) {
               LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error closing stream but continuing", e);
            }

            if (LOG.isDebugEnabled()) {
               LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "SAMLRequest is no valid XML Document");
            }
         } else {
            LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "SAMLRequest is no valid Base64");
         }
      } catch (SAML2Exception e) {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error parsing AuthnRequest: " + e.getMessage(), e);
      }

      LOG.debug("Invalid AuthnRequest: " + LogUtil.logSafeString(samlRequest));
      return Optional.empty();
   }
}