package com.airlock.iam.flow.domain.service.session;

import com.airlock.iam.flow.api.domain.model.session.FlowAttribute;
import com.airlock.iam.flow.api.domain.model.session.FlowAttributeRepository;
import com.google.common.base.Preconditions;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Stream;

public final class DefaultFlowAttributeRepository implements FlowAttributeRepository {
   private final Map<Class<? extends FlowAttribute<?>>, FlowAttribute<?>> attributes = new HashMap<>();

   @Override
   public <T extends FlowAttribute<?>> T putOrMergeAttribute(T value) {
      Preconditions.checkNotNull(value, "a flow session attribute must not be null");
      Class<T> keyType = (Class<T>)value.getRepositoryKey();
      T merged = (T)this.tryGetAttribute(keyType).map(old -> ((FlowAttribute<T>)old).mergeWith(value)).orElse(value);
      this.attributes.put(keyType, merged);
      return merged;
   }

   @Override
   public <T extends FlowAttribute<?>> Optional<T> tryGetAttribute(Class<T> key) {
      return Optional.ofNullable(this.nullableGet(key));
   }

   @Override
   public Stream<FlowAttribute<?>> values() {
      return this.attributes.values().stream();
   }

   @Override
   public <T extends FlowAttribute<?>> Optional<T> removeAttribute(Class<T> key) {
      return Optional.ofNullable((T)this.attributes.remove(key));
   }

   @Override
   public void removeAllAttributes() {
      this.attributes.clear();
   }

   private <T extends FlowAttribute<?>> T nullableGet(Class<T> key) {
      return key.cast(this.attributes.get(key));
   }
}