// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.vespa.athenz.identity;

import com.google.inject.Inject;
import com.yahoo.component.AbstractComponent;
import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
import com.yahoo.log.LogLevel;
import com.yahoo.vespa.athenz.api.AthenzService;
import com.yahoo.vespa.athenz.tls.KeyStoreType;
import com.yahoo.vespa.athenz.tls.SslContextBuilder;

import javax.net.ssl.SSLContext;
import java.io.File;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Logger;

/**
 * @author mortent
 * @author bjorncs
 */
public class SiaIdentityProvider extends AbstractComponent implements AthenzIdentityProvider {

    private static final Logger log = Logger.getLogger(SiaIdentityProvider.class.getName());

    private static final Duration REFRESH_INTERVAL = Duration.ofHours(1);

    private final AtomicReference<SSLContext> sslContext = new AtomicReference<SSLContext>();
    private final AthenzService service;
    private final File privateKeyFile;
    private final File certificateFile;
    private final File trustStoreFile;
    private final ScheduledExecutorService scheduler;

    @Inject
    public SiaIdentityProvider(SiaProviderConfig config) {
        this(new AthenzService(config.athenzDomain(), config.athenzService()),
             getPrivateKeyFile(config.keyPathPrefix(), config.athenzDomain(), config.athenzService()),
             getCertificateFile(config.keyPathPrefix(), config.athenzDomain(), config.athenzService()),
             new File(config.trustStorePath()),
             createScheduler());
    }

    public SiaIdentityProvider(AthenzService service,
                               Path siaPath,
                               File trustStoreFile) {
        this(service,
             getPrivateKeyFile(siaPath.toString(), service.getDomain().getName(), service.getName()),
             getCertificateFile(siaPath.toString(), service.getDomain().getName(), service.getName()),
             trustStoreFile,
             createScheduler());
    }

    public SiaIdentityProvider(AthenzService service,
                               File privateKeyFile,
                               File certificateFile,
                               File trustStoreFile,
                               ScheduledExecutorService scheduler) {
        this.service = service;
        this.privateKeyFile = privateKeyFile;
        this.certificateFile = certificateFile;
        this.trustStoreFile = trustStoreFile;
        this.scheduler = scheduler;
        this.sslContext.set(createIdentitySslContext());
        scheduler.scheduleAtFixedRate(this::reloadSslContext, REFRESH_INTERVAL.toMinutes(), REFRESH_INTERVAL.toMinutes(), TimeUnit.MINUTES);
    }

    private static ScheduledThreadPoolExecutor createScheduler() {
        return new ScheduledThreadPoolExecutor(1, runnable -> {
            Thread thread = new Thread(runnable);
            thread.setName("sia-identity-provider-sslcontext-updater");
            return thread;
        });
    }

    @Override
    public String getDomain() {
        return service.getDomain().getName();
    }

    @Override
    public String getService() {
        return service.getName();
    }

    @Override
    public SSLContext getIdentitySslContext() {
        return sslContext.get();
    }

    private SSLContext createIdentitySslContext() {
        return new SslContextBuilder()
                .withTrustStore(trustStoreFile, KeyStoreType.JKS)
                .withKeyStore(privateKeyFile, certificateFile)
                .build();
    }

    private void reloadSslContext() {
        log.log(LogLevel.DEBUG, "Updating SSLContext for identity " + service.getFullName());
        try {
            this.sslContext.set(createIdentitySslContext());
        } catch (Exception e) {
            log.log(LogLevel.SEVERE, "Failed to update SSLContext: " + e.getMessage(), e);
        }
    }

    private static File getCertificateFile(String rootPath, String domain, String service) {
        return Paths.get(rootPath, "certs", String.format("%s.%s.cert.pem", domain, service)).toFile();
    }

    private static File getPrivateKeyFile(String rootPath, String domain, String service) {
        return Paths.get(rootPath, "keys", String.format("%s.%s.key.pem", domain, service)).toFile();
    }

    @Override
    public void deconstruct() {
        try {
            scheduler.shutdownNow();
            scheduler.awaitTermination(90, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}

@@ -615,16 +615,89 @@
 ration;%0A
+import java.util.Set;%0Aimport java.util.concurrent.ConcurrentSkipListSet;%0A
 import j
@@ -738,24 +738,24 @@
 torService;%0A
-
 import java.
@@ -885,24 +885,60 @@
 cReference;%0A
+import java.util.function.Consumer;%0A
 import java.
@@ -1349,26 +1349,16 @@
 ference%3C
-SSLContext
 %3E();%0A   
@@ -1566,16 +1566,103 @@
 heduler;
+%0A    private final Set%3CConsumer%3CSSLContext%3E%3E listeners = new ConcurrentSkipListSet%3C%3E();
 %0A%0A    @I
@@ -3822,32 +3822,252 @@
 t.get();%0A    %7D%0A%0A
+    public void addReloadListener(Consumer%3CSSLContext%3E listener) %7B%0A        listeners.add(listener);%0A    %7D%0A%0A    public void removeReloadListener(Consumer%3CSSLContext%3E listener) %7B%0A        listeners.remove(listener);%0A    %7D%0A%0A
     private SSLC
@@ -4453,37 +4453,43 @@
 y %7B%0A            
-this.
+SSLContext 
 sslContext.set(c
@@ -4474,37 +4474,35 @@
 ntext sslContext
-.set(
+ = 
 createIdentitySs
@@ -4503,32 +4503,148 @@
 titySslContext()
+;%0A            this.sslContext.set(sslContext);%0A            listeners.forEach(listener -%3E listener.accept(sslContext)
 );%0A        %7D cat
@@ -5268,24 +5268,24 @@
 tdownNow();%0A
-
             
@@ -5330,24 +5330,55 @@
 t.SECONDS);%0A
+            listeners.clear();%0A
         %7D ca

Allow SiaIdentityProvider to notify listeners after reload