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

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.mapification.MapifiedValues;
import com.airlock.iam.base.domain.model.mapification.MapifiedValuesBuilder;
import com.airlock.iam.base.misc.json.JsonCoding;
import com.airlock.iam.saml2.domain.model.session.FlowFederationSession;
import com.sun.identity.saml2.assertion.impl.NameIDImpl;
import com.sun.identity.saml2.common.SAML2Exception;
import com.sun.identity.saml2.profile.IDPSession;
import com.sun.identity.saml2.profile.NameIDandSPpair;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

final class IDPSessionSerializer implements Serializer<IDPSession> {
   private static final String NAME_ID_AND_SPS_KEY = "nameIDandSPs";
   private static final String PENDING_LOGOUT_REQUEST_ID_KEY = "pendingLogoutRequestID";
   private static final String ORIGINATING_LOGOUT_REQUEST_ID_KEY = "originatingLogoutRequestID";
   private static final String ORIGINATING_LOGOUT_SP_ENTITY_ID_KEY = "originatingLogoutSPEntityID";
   private static final String META_ALIAS_KEY = "metaAlias";
   private static final String FED_SESSION = "fedSession";

   public String mapToStore(IDPSession session) {
      FailureException.checkTrue(
         session.getSession() == null || session.getSession() instanceof FlowFederationSession,
         () -> "Cannot handle session of type " + session.getSession().getClass()
      );
      Optional<FlowFederationSession> fedSession = Optional.ofNullable((FlowFederationSession)session.getSession());
      MapifiedValues values = MapifiedValuesBuilder.mapifiedValues()
         .putIfNotEmpty("nameIDandSPs", serializeNameIdAndSPairs(session.getNameIDandSPpairs()))
         .putIfPresent("pendingLogoutRequestID", Optional.ofNullable(session.getPendingLogoutRequestID()))
         .putIfPresent("originatingLogoutRequestID", Optional.ofNullable(session.getOriginatingLogoutRequestID()))
         .putIfPresent("originatingLogoutSPEntityID", Optional.ofNullable(session.getOriginatingLogoutSPEntityID()))
         .putIfPresent("metaAlias", Optional.ofNullable(session.getMetaAlias()))
         .putIfPresent("fedSession", fedSession.map(FlowFederationSession::mapify))
         .build();
      return JsonCoding.toJson(values.asMap());
   }

   private static List<List<String>> serializeNameIdAndSPairs(List<NameIDandSPpair> nameIDandSPpairs) {
      return nameIDandSPpairs.stream().map(nsp -> {
         try {
            return List.of(nsp.getNameID().toXMLString(true, true), nsp.getSPEntityID());
         } catch (SAML2Exception e) {
            throw new FailureException(e);
         }
      }).toList();
   }

   public IDPSession mapFromStore(String serialized) {
      MapifiedValues values = new MapifiedValues(JsonCoding.fromJson(serialized));
      FlowFederationSession fedSession = values.demapifyOptional("fedSession", FlowFederationSession::demapify).orElse(null);
      IDPSession session = new IDPSession(fedSession, deserializeNameIdAndSPairs(values.getList("nameIDandSPs")));
      values.<String>getOptional("pendingLogoutRequestID").ifPresent(session::setPendingLogoutRequestID);
      values.<String>getOptional("originatingLogoutRequestID").ifPresent(session::setOriginatingLogoutRequestID);
      values.<String>getOptional("originatingLogoutSPEntityID").ifPresent(session::setOriginatingLogoutSPEntityID);
      values.<String>getOptional("metaAlias").ifPresent(session::setMetaAlias);
      return session;
   }

   private static List<NameIDandSPpair> deserializeNameIdAndSPairs(Object object) {
      List<List<String>> list = (List<List<String>>)object;
      return list.stream().map(IDPSessionSerializer::deserializeNameIdAndSPair).collect(Collectors.toCollection(ArrayList::new));
   }

   private static NameIDandSPpair deserializeNameIdAndSPair(List<String> pair) {
      try {
         return new NameIDandSPpair(new NameIDImpl(pair.getFirst()), pair.get(1));
      } catch (SAML2Exception e) {
         throw new FailureException(e);
      }
   }
}