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

import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.saml2.domain.model.assertion.FlowSaml2Attribute;
import com.airlock.iam.saml2.domain.model.logging.Saml2LogTypes;
import com.airlock.iam.saml2.domain.model.session.FlowFederationSession;
import com.google.common.annotations.VisibleForTesting;
import com.sun.identity.plugin.session.SessionException;
import com.sun.identity.plugin.session.SessionManager;
import com.sun.identity.plugin.session.SessionProvider;
import com.sun.identity.saml2.assertion.Attribute;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.common.SAML2Utils;
import com.sun.identity.saml2.plugins.IDPAttributeMapper;
import java.util.List;

public class IDPAttributeMapperImpl implements IDPAttributeMapper {
   private static final ELogger LOG = ELoggerFactory.getLogger(IDPAttributeMapperImpl.class);

   @Override
   public List<Attribute> getAttributes(Object session, String localEntityID, String remoteEntityID, String realm) throws SAML2Exception {
      if (localEntityID == null) {
         throw new SAML2Exception("Local entity ID empty");
      }

      if (realm == null) {
         throw new SAML2Exception("Realm empty");
      }

      SessionProvider sessionProvider;
      try {
         sessionProvider = SessionManager.getProvider();
      } catch (SessionException e) {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error getting SessionManager", e);
         throw new SAML2Exception(e);
      }

      return this.getAttributes(session, sessionProvider);
   }

   @VisibleForTesting
   List<Attribute> getAttributes(Object session, SessionProvider sessionProvider) throws SAML2Exception {
      if (session == null) {
         throw new SAML2Exception("Session is null");
      }

      try {
         if (!sessionProvider.isValid(session)) {
            LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Invalid Session");
            throw new SAML2Exception("Invalid Session");
         }
      } catch (SessionException e) {
         LOG.warn(Saml2LogTypes.FEDERATION_LIBRARY_MESSAGE, "Error accessing Session", e);
         throw new SAML2Exception(e);
      }

      FlowFederationSession federationSession = FlowFederationSession.toFlowFederationSessionOrThrow(
         session, () -> new SAML2Exception("Session is not of correct type: " + session.getClass().getName())
      );
      return federationSession.getPropagationAttributes()
         .orElseThrow(() -> new SAML2Exception("Session does not contain propagation attributes"))
         .getAssertionAttributes()
         .stream()
         .map(this::createSamlAttribute)
         .toList();
   }

   private Attribute createSamlAttribute(FlowSaml2Attribute attr) {
      try {
         Attribute samlAttribute = SAML2Utils.getSAMLAttribute(attr.getAttributeName(), attr.getAttributeValues().toArray(String[]::new));
         samlAttribute.setNameFormat(attr.getNameFormat());
         return samlAttribute;
      } catch (SAML2Exception e) {
         throw new IllegalArgumentException("Error creating SAML2 Attribute: " + e.getMessage(), e);
      }
   }
}