package com.airlock.iam.oauth2.domain.model.valueprovider;

import com.airlock.iam.base.api.domain.model.plugin.LogicForConfig;
import com.airlock.iam.flow.shared.api.domain.model.valueprovider.ValueMapProvider;
import com.airlock.iam.flow.shared.api.domain.service.FlowAttributeService;
import com.airlock.iam.oauth2.application.configuration.valueprovider.OAuth2TokensValueMapProviderConfig;
import com.airlock.iam.oauth2.domain.model.oauth2.session.OAuth2CredentialAuthenticationFlowSessionAttribute;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMap.Builder;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.util.Collections;
import java.util.Map;

@LogicForConfig(OAuth2TokensValueMapProviderConfig.class)
public class OAuth2CredentialTokensValueMapProvider implements ValueMapProvider {
   static final String ACCESS_TOKEN_KEY = "access_token";
   static final String ID_TOKEN_KEY = "id_token";
   private final Provider<FlowAttributeService> attributes;

   @Inject
   public OAuth2CredentialTokensValueMapProvider(Provider<FlowAttributeService> attributes) {
      this.attributes = attributes;
   }

   @Override
   public Map<String, Object> provideValueMap() {
      return ((FlowAttributeService)this.attributes.get()).findAttribute(OAuth2CredentialAuthenticationFlowSessionAttribute.class).map(attribute -> {
         Builder<String, Object> map = ImmutableMap.builder().put("access_token", attribute.getAccessToken().value());
         attribute.getIdToken().ifPresent(idToken -> map.put("id_token", idToken.tokenString()));
         return map.build();
      }).orElse(Collections.emptyMap());
   }
}