package com.airlock.iam.common.infrastructure.service;

import com.airlock.iam.common.application.service.JwksService;
import com.airlock.iam.common.domain.model.jwk.JsonWebKeySet;
import com.airlock.iam.common.domain.model.jwt.SignatureVerificationInformation;
import com.airlock.iam.common.domain.service.JwksPublicKeyRegistry;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import java.util.Collection;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Objects;
import java.util.Set;

public class JwksServiceImpl implements JwksService {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwksServiceImpl.class);
   private final Set<JwksPublicKeyRegistry> registries = Collections.newSetFromMap(new IdentityHashMap<>());
   private JsonWebKeySet cache;

   @Override
   public void register(JwksPublicKeyRegistry registry) {
      this.registries.add(registry);
   }

   @Override
   public void deregister(JwksPublicKeyRegistry registry) {
      this.registries.remove(registry);
   }

   @Override
   public void reload() {
      this.cache = JsonWebKeySet.fromJson(new JWKSet(this.loadKeys()).toString(true)).orElseThrow();
   }

   private List<JWK> loadKeys() {
      return this.registries
         .stream()
         .map(JwksServiceImpl::safelyGetSigningCertificates)
         .map(JWKSetExtractor::new)
         .map(JwksServiceImpl::safelyGetJwkSet)
         .filter(Objects::nonNull)
         .map(JWKSet::getKeys)
         .flatMap(Collection::stream)
         .distinct()
         .toList();
   }

   private static List<SignatureVerificationInformation> safelyGetSigningCertificates(JwksPublicKeyRegistry registry) {
      try {
         return registry.getSigningCertificatesForJwks();
      } catch (Exception e) {
         LOG.info("Failed to get signing certificate for jwks", e);
         return List.of();
      }
   }

   private static JWKSet safelyGetJwkSet(JWKSetExtractor extractor) {
      try {
         return extractor.getJwkSet();
      } catch (Exception e) {
         LOG.info("Error retrieving keys from registry: " + e.getMessage());
         return null;
      }
   }

   @Override
   public JsonWebKeySet jsonWebKeySet() {
      if (this.cache == null) {
         this.reload();
      }

      return this.cache;
   }
}