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

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.saml2.application.configuration.SamlFederationConfig;
import com.sun.identity.saml2.plugins.DefaultLibrarySPAccountMapper;
import com.sun.identity.saml2.plugins.DefaultSPAttributeMapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.function.Function;
import org.w3c.dom.Document;

public class ExtendedSpMetadataTransformer implements Function<String, String> {
   private final SamlFederationConfig samlFederationConfig;

   public ExtendedSpMetadataTransformer(SamlFederationConfig samlFederationConfig) {
      this.samlFederationConfig = samlFederationConfig;
   }

   public String apply(String xml) {
      try (InputStream is = new ByteArrayInputStream(xml.getBytes(StandardCharsets.UTF_8))) {
         Document doc = MetadataUtils.loadDocument(is);
         MetadataUtils.setExtendedAttributeValue(doc, "spAdapter", this.samlFederationConfig.getSpAdapter());
         MetadataUtils.setExtendedAttributeValue(doc, "spAccountMapper", DefaultLibrarySPAccountMapper.class.getName());
         MetadataUtils.setExtendedAttributeValue(doc, "spAttributeMapper", DefaultSPAttributeMapper.class.getName());
         MetadataUtils.setExtendedAttributeValue(doc, "spAuthncontextMapper", this.samlFederationConfig.getSpAuthnContextMapper());
         return new String(MetadataUtils.getDocumentBytes(doc), StandardCharsets.UTF_8);
      } catch (IOException e) {
         throw new FailureException("Error closing ByteArrayInputStream: " + e.getMessage());
      }
   }
}