package com.atlassian.bitbucket.mesh.management;

import com.atlassian.bitbucket.mesh.Product;
import com.atlassian.bitbucket.mesh.annotation.NotProfiled;
import com.atlassian.bitbucket.mesh.annotation.Profiled;
import com.atlassian.bitbucket.mesh.app.BuildInfo;
import com.atlassian.bitbucket.mesh.auth.*;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.config.ConfigurationUpdatedEvent;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.git.builder.version.VersionStdioHandler;
import com.atlassian.bitbucket.mesh.io.FilePermission;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.io.SetFilePermissionRequest;
import com.atlassian.bitbucket.mesh.layout.StorageLayout;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.UpdateTopologyParameters;
import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.atlassian.bitbucket.mesh.repository.MarkInconsistentParameters;
import com.atlassian.bitbucket.mesh.repository.ReplicaStateRegistry;
import com.atlassian.bitbucket.mesh.throttle.ThrottleService;
import com.atlassian.bitbucket.mesh.util.PropertiesUtils;
import com.atlassian.bitbucket.mesh.util.Version;
import com.google.common.collect.ImmutableMap;
import org.apache.commons.io.file.SimplePathVisitor;
import org.apache.commons.io.output.CloseShieldOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.env.PropertySourceLoader;
import org.springframework.core.env.ConfigurableEnvironment;
import org.springframework.core.env.EnumerablePropertySource;
import org.springframework.core.env.PropertySource;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PropertiesLoaderUtils;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.channels.OverlappingFileLockException;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.BasicFileAttributes;
import java.security.PublicKey;
import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import static com.atlassian.bitbucket.mesh.MeshConstants.PROFILES_CLOUD;
import static com.atlassian.bitbucket.mesh.MeshConstants.PROFILES_SIDECAR;
import static com.atlassian.bitbucket.mesh.management.SupportUtils.APPLICATION_PROPERTIES_JSON;
import static com.atlassian.bitbucket.mesh.repository.ManagedRepository.Id.parseId;
import static com.atlassian.bitbucket.mesh.rpc.ManagedConfigurationProperties.NODE_ID;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toMap;
import static org.apache.commons.lang3.StringUtils.isNotBlank;

@Profiled
public class DefaultNodeManager implements NodeManager {

    static final String CONFIG_BANNER = "# This configuration file is managed by " + Product.FULL_NAME + ".\n" +
            "# Centrally managed settings may be overwritten by the system at any time.\n\n";

    private static final Logger log = LoggerFactory.getLogger(DefaultNodeManager.class);

    private final AuthenticationManager authenticationManager;
    private final BuildInfo buildInfo;
    private final GitProcessBuilderFactory builderFactory;
    private final ApplicationConfig config;
    private final Path configFile;
    private final CountDownLatch configLatch;
    private final Object configLock;
    private final Properties defaultProperties;
    private final ConfigurableEnvironment environment;
    private final EventPublisher eventPublisher;
    private final FileSanitizer fileSanitizer;
    private final MeshManager meshManager;
    private final PropertySourceLoader propertySourceLoader;
    private final ReplicaStateRegistry replicaStateRegistry;
    private final StorageLayout storageLayout;
    private final SupportInfoRegistry supportInfoRegistry;
    private final Path tempDir;
    private final ThrottleService throttleService;
    private final CountDownLatch topologyLatch;

    private volatile boolean configured;

    public DefaultNodeManager(AuthenticationManager authenticationManager, BuildInfo buildInfo,
                              GitProcessBuilderFactory builderFactory, ApplicationConfig config,
                              ConfigurableEnvironment environment, EventPublisher eventPublisher,
                              FileSanitizer fileSanitizer, MeshManager meshManager,
                              PropertySourceLoader propertySourceLoader, ReplicaStateRegistry replicaStateRegistry,
                              StorageLayout storageLayout, SupportInfoRegistry supportInfoRegistry,
                              ThrottleService throttleService, boolean blockForTopology) {
        this.authenticationManager = authenticationManager;
        this.buildInfo = buildInfo;
        this.builderFactory = builderFactory;
        this.config = config;
        this.environment = environment;
        this.eventPublisher = eventPublisher;
        this.fileSanitizer = fileSanitizer;
        this.meshManager = meshManager;
        this.propertySourceLoader = propertySourceLoader;
        this.replicaStateRegistry = replicaStateRegistry;
        this.storageLayout = storageLayout;
        this.supportInfoRegistry = supportInfoRegistry;
        this.throttleService = throttleService;

        configFile = storageLayout.getConfigFile();
        configLatch = new CountDownLatch(1);
        configLock = new Object();
        tempDir = storageLayout.getTempDir();
        topologyLatch = new CountDownLatch(blockForTopology ? 1 : 0);
        defaultProperties = loadDefaultConfig(environment);
    }

    @Override
    public boolean awaitConfiguration(@Nonnull Duration timeout) throws InterruptedException {
        requireNonNull(timeout, "timeout");

        if (configured) {
            return true;
        }
        Instant deadline = Instant.now().plus(timeout);
        if (configLatch.await(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
            if (topologyLatch.await(Duration.between(Instant.now(), deadline).toMillis(), TimeUnit.MILLISECONDS)) {
                //This method is likely to be called extensively, so we lazily "cache" a true result after the
                //first call gets a true result (since the value can never go false again)
                configured = true;
                return true;
            }
        }
        return false;
    }

    @Nonnull
    @Override
    public Path createSupportZip(@Nonnull CreateSupportZipParameters parameters) {
        requireNonNull(parameters, "parameters");

        Path supportZip = null;
        try {
            supportZip = Files.createTempFile(tempDir, "mesh-support-", ".zip");

            try (ZipOutputStream zipStream = new ZipOutputStream(Files.newOutputStream(supportZip))) {
                zipStream.setLevel(9);
                addApplicationPropertiesToSupportZip(zipStream);
                if (parameters.isWithLogs()) {
                    Path logDir = storageLayout.getLogDir();
                    addLogsToSupportZip(zipStream, logDir);
                }

                if (parameters.isWithConfig()) {
                    addConfigFileToSupportZip(zipStream, storageLayout.getConfigFile());
                    addConfigFileToSupportZip(zipStream, storageLayout.getNodeConfigFile());
                }
            }

            return supportZip;
        } catch (IOException e) {
            if (supportZip != null) {
                MoreFiles.deleteQuietly(supportZip);
            }
            log.error("Could not successfully construct the support zip", e);
            throw new SupportZipFailedException("Support zip could not be created", e);
        }
    }

    @Nonnull
    @Override
    public Map<String, String> getConfigurationProperties() {
        Set<String> keys = environment.getPropertySources().stream()
                .filter(ps -> ps instanceof EnumerablePropertySource)
                .flatMap(ps -> Arrays.stream(EnumerablePropertySource.class.cast(ps).getPropertyNames()))
                .filter(key -> defaultProperties.containsKey(key) || key.startsWith("logging.logger."))
                .collect(Collectors.toSet());

        ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
        for (String key : keys) {
            String value = environment.getProperty(key);
            if (isNotBlank(value)) {
                String defaultValue = defaultProperties.getProperty(key);
                if (!value.equals(defaultValue) &&
                        (defaultValue == null || !value.equals(environment.resolvePlaceholders(defaultValue)))) {
                    builder.put(key, value);
                }
            }
        }
        return builder.build();
    }

    @Nonnull
    @Override
    @NotProfiled
    public String getId() {
        return config.getNode().getId();
    }

    @Nonnull
    @Override
    @NotProfiled
    public Map<String, String> getNodeInfo() {
        return ImmutableMap.of(
                "id", getId(),
                "buildVersion", buildInfo.getBuildVersion());
    }

    @Nonnull
    @Override
    public SupportInfo getSupportInfo() {
        Version currentVersion = builderFactory.builder()
                .version()
                .build(new VersionStdioHandler())
                .run();
        Version startupVersion = config.getGit().getBinary().getVersion();

        return new SupportInfo(new GitInfo(currentVersion, startupVersion), Arrays.stream(throttleService.getSummaries())
                .map(TicketInfo::new)
                .collect(toImmutableList()));
    }

    @Override
    @NotProfiled
    public boolean isConfigured() {
        return configured || (configLatch.getCount() == 0L && topologyLatch.getCount() == 0L);
    }

    @Override
    public PublicKey provision(@Nonnull ProvisionParameters parameters) {
        synchronized (configLock) {
            if (authenticationManager != null) {
                //A node can only be provisioned anonymously if not already provisioned
                //(see AuthenticationServerInterceptor), or re-provisioned using a control plane key.
                ApplicationRole role = authenticationManager.getPrincipal().map(Principal::getRole).orElse(null);
                if (role != null && role != ApplicationRole.CONTROL_PLANE) {
                    log.warn("Blocked a (re-)provisioning attempt by {}",
                            authenticationManager.getPrincipal().get().getId());
                    throw new AuthorisationException("A node can only be (re-)provisioned by the control plane");
                }
            }

            String nodeId = parameters.getNodeId();
            if (nodeId != null) {
                try {
                    if (updateConfigFile(ImmutableMap.of(NODE_ID, nodeId))) {
                        //Refresh the Spring property source to make the updated properties visible
                        refreshPropertySource();
                        //And reload the config
                        config.reload();
                    }
                } catch (IOException e) {
                    log.error("Failed to update node ID to {}", nodeId, e);
                    throw new ConfigurationFailedException("Configuration was not applied", e);
                }
            }

            if (authenticationManager == null) {
                log.info("Not provisioning; authentication is disabled");
                eventPublisher.publish(new NodeProvisionedEvent());
                return null;
            }

            PublicKey result = authenticationManager.provision(parameters.getPublicKey());
            eventPublisher.publish(new NodeProvisionedEvent());
            return result;
        }
    }

    @Override
    public void updateConfiguration(@Nonnull UpdateConfigurationParameters parameters) {
        requireNonNull(parameters, "parameters");

        synchronized (configLock) {
            Map<String, String> properties = parameters.getProperties();
            if (properties != null) {
                try {
                    //Update the mesh.properties file with the new properties
                    if (updateConfigFile(properties)) {
                        //Refresh the Spring property source to make the updated properties visible
                        refreshPropertySource();
                        //And reload the config
                        config.reload();
                    }

                    //Only unlatch if configuration was received from the control plane, was saved and, if necessary,
                    //refreshed/reloaded. If this happens multiple times, all times after the first will be no-ops
                    if (parameters.isFromControlPlane()) {
                        configLatch.countDown();
                    }
                } catch (IOException e) {
                    log.error("Failed to (fully) apply configuration update: {}", parameters.getProperties(), e);
                    throw new ConfigurationFailedException("Configuration was not applied", e);
                }
            }

            UpdatePublicKeysParameters publicKeysParameters = parameters.getPublicKeys();
            if (publicKeysParameters != null) {
                if (authenticationManager == null) {
                    log.debug("Not updating public keys; authentication is disabled");
                } else {
                    try {
                        authenticationManager.updatePublicKeys(publicKeysParameters);
                    } catch (RuntimeException e) {
                        log.error("Failed to update public keys", e);
                        throw new ConfigurationFailedException("Public keys were not applied", e);
                    }
                }
            }

            UpdateTopologyParameters topologyParameters = parameters.getTopology();
            if (topologyParameters != null) {
                try {
                    meshManager.updateTopology(topologyParameters);
                    if (parameters.isFromControlPlane()) {
                        topologyLatch.countDown();
                    }
                } catch (RuntimeException e) {
                    log.error("Failed to update topology", e);
                    throw new ConfigurationFailedException("Topology was not applied", e);
                }
            }

            try {
                //Filter inconsistent replicas for partitions that are not hosted by this node. This
                //prevents repair attempts that can't succeed because the partition doesn't exist
                //on this node yet.
                Map<String, Long> inconsistentReplicas = parameters.getInconsistentReplicas().entrySet().stream()
                        .filter(entry ->
                                parseId(entry.getKey())
                                        .map(ManagedRepository.Id::getPartition)
                                        .map(meshManager::getPartition)
                                        .isPresent())
                        .collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
                replicaStateRegistry.markInconsistent(new MarkInconsistentParameters.Builder()
                        .complete(true)
                        .fromControlPlane(parameters.isFromControlPlane())
                        .inconsistentReplicas(inconsistentReplicas)
                        .build());
            } catch (RuntimeException e) {
                log.error("Failed to update inconsistent replica registry", e);
                throw new ConfigurationFailedException("Inconsistent replica registrations were not applied", e);
            }
        }
    }

    private static void addLogsToSupportZip(ZipOutputStream zipStream, Path logDir) throws IOException {
        Files.walkFileTree(logDir, new SimplePathVisitor() {

            @Override
            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                Path targetFile = logDir.relativize(file);

                zipStream.putNextEntry(new ZipEntry("log/" + targetFile));
                try {
                    Files.copy(file, zipStream);
                } catch (IOException e) {
                    log.error("Could not successfully add the log file {} to the support zip", targetFile);
                    throw e;
                } finally {
                    zipStream.closeEntry();
                }

                return FileVisitResult.CONTINUE;
            }
        });
    }

    private static boolean isSensitive(String key) {
        return "authentication.token".equals(key);
    }

    private static Properties loadDefaultConfig(ConfigurableEnvironment environment) {
        List<Resource> configResources = new ArrayList<>(3);
        configResources.add(new ClassPathResource("application-private.properties"));
        configResources.add(new ClassPathResource("application-public.properties"));
        if (environment.acceptsProfiles(PROFILES_CLOUD)) {
            //In Cloud mode, mesh-cloud.properties contains additional defaults
            configResources.add(new ClassPathResource("mesh-cloud.properties"));
        } else if (environment.acceptsProfiles(PROFILES_SIDECAR)) {
            //In sidecar mode, mesh-sidecar.properties contains additional defaults
            configResources.add(new ClassPathResource("mesh-sidecar.properties"));
        }

        Properties properties = new Properties();
        for (Resource resource : configResources) {
            try {
                PropertiesLoaderUtils.fillProperties(properties, resource);
            } catch (IOException e) {
                log.warn("Problem loading properties from {}", resource.getFilename(), e);
            }
        }

        return properties;
    }

    private void addApplicationPropertiesToSupportZip(ZipOutputStream zipStream) throws IOException {
        zipStream.putNextEntry(new ZipEntry(APPLICATION_PROPERTIES_JSON));
        try {
            supportInfoRegistry.writeSupportInfo(CloseShieldOutputStream.wrap(zipStream));
        } catch (IOException e) {
            log.error("Could not successfully add {} to the support zip", APPLICATION_PROPERTIES_JSON);
            throw e;
        } finally {
            zipStream.closeEntry();
        }
    }

    private void addConfigFileToSupportZip(ZipOutputStream zipStream, Path configFile) throws IOException {
        if (Files.exists(configFile)) {
            Path sanitizedFile = fileSanitizer.sanitize(configFile);

            zipStream.putNextEntry(new ZipEntry("config/" + configFile.getFileName()));
            try {
                Files.copy(sanitizedFile, zipStream);
            } catch (IOException e) {
                log.error("Could not successfully add the sanitized file {} to the support zip", configFile);
                throw e;
            } finally {
                zipStream.closeEntry();
                MoreFiles.deleteQuietly(sanitizedFile);
            }
        }
    }

    private PropertySource<?> loadConfigPropertySource(String name) throws IOException {
        List<PropertySource<?>> sources = propertySourceLoader.load(name, new FileSystemResource(configFile));
        return sources.isEmpty() ? null : sources.iterator().next();
    }

    private boolean performConfigFileUpdate(Map<String, String> propertyMap, Properties properties) throws IOException {
        Set<String> updatedKeys = PropertiesUtils.update(properties, configFile, CONFIG_BANNER);
        MoreFiles.setPermissions(new SetFilePermissionRequest.Builder(configFile)
                .ownerPermissions(EnumSet.of(FilePermission.READ, FilePermission.WRITE))
                .build());

        if (updatedKeys.isEmpty()) {
            return false;
        }

        eventPublisher.publish(new ConfigurationUpdatedEvent(updatedKeys));
        if (log.isDebugEnabled()) {
            StringBuilder message = new StringBuilder("Updated properties: ");
            propertyMap.forEach((key, value) -> message.append("\n\t")
                    .append(updatedKeys.contains(key) ? "* " : "  ")
                    .append(key)
                    .append('=')
                    .append(isSensitive(key) ? "*****************" : value));
            log.debug(message.toString());
        }
        return true;
    }

    private void refreshPropertySource() throws IOException {
        String filename = configFile.getFileName().toString();
        String sourceName = environment.getPropertySources().stream()
                .map(PropertySource::getName)
                .filter(name -> name.contains(filename))
                .findFirst()
                .orElse(null);

        if (sourceName == null) {
            //No property source was found for the config file. This means that the config file did not exist
            PropertySource<?> updated = loadConfigPropertySource("applicationConfig: [file:" + filename + "]");
            if (updated != null) {
                try {
                    environment.getPropertySources().addAfter("random", updated);
                } catch (RuntimeException e) {
                    //"random" PropertySource was not defined. This happens in the tests where the full spring boot
                    //environment isn't set up. Add the property source after systemEnvironment instead.
                    environment.getPropertySources().addAfter("systemEnvironment", updated);
                }
            }
        } else {
            PropertySource<?> updated = loadConfigPropertySource(sourceName);
            if (updated == null) {
                //After the update, the config file is empty. Remove the existing property source
                environment.getPropertySources().remove(sourceName);
            } else {
                //Otherwise, replace the existing PropertySource with the updated one
                environment.getPropertySources().replace(sourceName, updated);
            }
        }
    }

    private boolean updateConfigFile(Map<String, String> propertyMap) throws IOException {
        if (propertyMap.isEmpty()) {
            //If there's no configuration to apply, short-circuit processing
            return false;
        }

        //Make sure the directory exists
        MoreFiles.mkdir(configFile.getParent());
        Properties properties = new Properties();
        propertyMap.forEach((key, value) -> {
            if (key.startsWith("logging.logger.")) {
                properties.put(key, value);
                return;
            }
            String defaultValue = defaultProperties.getProperty(key);
            if (defaultValue == null) {
                log.info("Ignoring unknown config property {}", key);
                return;
            }
            if (Objects.equals(value, defaultValue) ||
                    Objects.equals(value, environment.resolvePlaceholders(defaultValue))) {
                //The new value is equal to the defaultValue. Set the value to an empty string in properties, which
                //will trigger removal of the property from the config file.
                properties.put(key, "");
            } else {
                //Non-default value set, write this to the config file
                properties.put(key, value);
            }
        });

        try {
            return performConfigFileUpdate(propertyMap, properties);
        } catch (OverlappingFileLockException e) {
            //Could not obtain a lock on the config file
            Path lockFile = configFile.resolveSibling(configFile.getFileName() + ".lock");
            if (MoreFiles.getDurationSinceLastModified(lockFile).toMinutes() > 15) {
                //The lock file is older than 15 minutes, assume it's stale and delete it
                if (MoreFiles.deleteQuietly(lockFile)) {
                    return performConfigFileUpdate(propertyMap, properties);
                }
            }
            return false;
        }
    }
}
