/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.tools;

import org.apache.kafka.common.utils.AppInfoParser;
import org.apache.kafka.connect.runtime.isolation.PluginSource;
import org.apache.kafka.connect.runtime.isolation.PluginType;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.PrintStream;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileSystem;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
import java.util.Collection;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.stream.Stream;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

/**
 * An in-memory workspace for manipulating {@link java.util.ServiceLoader} manifest files.
 * <p>Use {@link #forSource(PluginSource)} to get a workspace scoped to a single plugin location, which is able
 * to accept simulated reads and writes of manifests.
 * Write the simulated changes to disk via {@link #commit(boolean)}.
 */
public class ManifestWorkspace {

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

    private static final String MANIFEST_PREFIX = "META-INF/services/";
    private static final Path MANAGED_PATH = Paths.get("connect-plugin-path-shim-1.0.0.jar");
    private static final String MANIFEST_HEADER = "# Generated by connect-plugin-path.sh " + AppInfoParser.getVersion();
    private final PrintStream out;
    private final List<SourceWorkspace<?>> workspaces;
    private final Map<Path, Path> temporaryOverlayFiles;

    public ManifestWorkspace(PrintStream out) {
        this.out = out;
        workspaces = new ArrayList<>();
        temporaryOverlayFiles = new HashMap<>();
    }

    public SourceWorkspace<?> forSource(PluginSource source) throws IOException {
        SourceWorkspace<?> sourceWorkspace = switch (source.type()) {
            case CLASSPATH -> new ClasspathWorkspace(source);
            case MULTI_JAR -> new MultiJarWorkspace(source);
            case SINGLE_JAR -> new SingleJarWorkspace(source);
            case CLASS_HIERARCHY -> new ClassHierarchyWorkspace(source);
            default -> throw new IllegalStateException("Unknown source type " + source.type());
        };
        workspaces.add(sourceWorkspace);
        return sourceWorkspace;
    }

    /**
     * Commits all queued changes to disk
     * @return true if any workspace wrote changes to disk, false if all workspaces did not have writes to apply
     * @throws IOException if an error occurs reading or writing to the filesystem
     * @throws TerseException if a path is not writable on disk and should be.
     */
    public boolean commit(boolean dryRun) throws IOException, TerseException {
        boolean changed = false;
        for (SourceWorkspace<?> workspace : workspaces) {
            changed |= workspace.commit(dryRun);
        }
        return changed;
    }

    /**
     * A workspace scoped to a single plugin source.
     * <p>Buffers simulated reads and writes to the plugin path before they can be written to disk.
     * @param <T> The data structure used by the workspace to store in-memory manifests internally.
     */
    public abstract static class SourceWorkspace<T> {
        private final Path location;
        private final PluginSource.Type type;
        protected final T initial;
        protected final T manifests;

        private SourceWorkspace(PluginSource source) throws IOException {
            this.location = source.location();
            this.type = source.type();
            this.initial = load(source);
            this.manifests = load(source);
        }

        public Path location() {
            return location;
        }

        public PluginSource.Type type() {
            return type;
        }

        protected abstract T load(PluginSource source) throws IOException;

        public abstract boolean hasManifest(PluginType type, String className);

        public abstract void forEach(BiConsumer<String, PluginType> consumer);

        public abstract void addManifest(PluginType type, String pluginClass);

        public abstract void removeManifest(PluginType type, String pluginClass);

        protected abstract boolean commit(boolean dryRun) throws TerseException, IOException;

        protected static Map<PluginType, Set<String>> loadManifest(URL baseUrl) throws MalformedURLException {
            Map<PluginType, Set<String>> manifests = new EnumMap<>(PluginType.class);
            for (PluginType type : PluginType.values()) {
                Set<String> result;
                try {
                    URL u = new URL(baseUrl, MANIFEST_PREFIX + type.superClass().getName());
                    result = parse(u);
                } catch (RuntimeException e) {
                    result = new LinkedHashSet<>();
                }
                manifests.put(type, result);
            }
            return manifests;
        }

        protected static URL jarBaseUrl(URL fileUrl) throws MalformedURLException {
            return new URL("jar", "", -1, fileUrl + "!/", null);
        }

        protected static void forEach(Map<PluginType, Set<String>> manifests, BiConsumer<String, PluginType> consumer) {
            manifests.forEach((type, classNames) -> classNames.forEach(className -> consumer.accept(className, type)));
        }
    }

    /**
     * A single jar can only contain one manifest per plugin type.
     */
    private class SingleJarWorkspace extends SourceWorkspace<Map<PluginType, Set<String>>> {

        private SingleJarWorkspace(PluginSource source) throws IOException {
            super(source);
            assert source.urls().length == 1;
        }

        @Override
        protected Map<PluginType, Set<String>> load(PluginSource source) throws IOException {
            return loadManifest(jarBaseUrl(source.urls()[0]));
        }

        @Override
        public boolean hasManifest(PluginType type, String className) {
            return manifests.get(type).contains(className);
        }

        @Override
        public void forEach(BiConsumer<String, PluginType> consumer) {
            forEach(manifests, consumer);
        }

        @Override
        public void addManifest(PluginType type, String pluginClass) {
            manifests.get(type).add(pluginClass);
        }

        @Override
        public void removeManifest(PluginType type, String pluginClass) {
            manifests.get(type).remove(pluginClass);
        }

        @Override
        protected boolean commit(boolean dryRun) throws IOException, TerseException {
            if (startSync(dryRun, location(), initial, manifests)) {
                rewriteJar(dryRun, location(), manifests);
                return true;
            }
            return false;
        }
    }

    /**
     * A classpath workspace is backed by multiple jars, and is not writable.
     * The in-memory format is a map from jar path to the manifests contained in that jar.
     * The control flow of the caller should not perform writes, so these exceptions indicate a bug in the program.
     */
    private class ClasspathWorkspace extends SourceWorkspace<Map<Path, Map<PluginType, Set<String>>>> {

        private ClasspathWorkspace(PluginSource source) throws IOException {
            super(source);
        }

        @Override
        protected Map<Path, Map<PluginType, Set<String>>> load(PluginSource source) throws IOException {
            Map<Path, Map<PluginType, Set<String>>> manifestsBySubLocation = new HashMap<>();
            for (URL url : source.urls()) {
                Path jarPath = Paths.get(url.getPath());
                manifestsBySubLocation.put(jarPath, loadManifest(jarBaseUrl(url)));
            }
            return manifestsBySubLocation;
        }

        public boolean hasManifest(PluginType type, String className) {
            return manifests.values()
                    .stream()
                    .map(m -> m.get(type))
                    .anyMatch(s -> s.contains(className));
        }

        public void forEach(BiConsumer<String, PluginType> consumer) {
            manifests.values().forEach(m -> forEach(m, consumer));
        }

        @Override
        public void addManifest(PluginType type, String pluginClass) {
            throw new UnsupportedOperationException("Cannot change the contents of the classpath");
        }

        @Override
        public void removeManifest(PluginType type, String pluginClass) {
            throw new UnsupportedOperationException("Cannot change the contents of the classpath");
        }

        @Override
        protected boolean commit(boolean dryRun) throws IOException, TerseException {
            // There is never anything to commit for the classpath
            return false;
        }
    }

    /**
     * A multi-jar workspace is similar to the classpath workspace because it has multiple jars.
     * However, the multi-jar workspace is writable, and injects a managed jar where it writes added manifests.
     */
    private class MultiJarWorkspace extends ClasspathWorkspace {

        private MultiJarWorkspace(PluginSource source) throws IOException {
            super(source);
        }

        @Override
        protected Map<Path, Map<PluginType, Set<String>>> load(PluginSource source) throws IOException {
            Map<Path, Map<PluginType, Set<String>>> manifests = super.load(source);
            // In addition to the normal multi-jar paths, inject a managed jar where we can add manifests.
            Path managedPath = source.location().resolve(MANAGED_PATH);
            URL url = managedPath.toUri().toURL();
            manifests.put(managedPath, loadManifest(jarBaseUrl(url)));
            return manifests;
        }

        @Override
        public void addManifest(PluginType type, String pluginClass) {
            // Add plugins to the managed manifest
            manifests.get(location().resolve(MANAGED_PATH)).get(type).add(pluginClass);
        }

        @Override
        public void removeManifest(PluginType type, String pluginClass) {
            // If a plugin appears in multiple manifests, remove it from all of them.
            for (Map<PluginType, Set<String>> manifestState : manifests.values()) {
                manifestState.get(type).remove(pluginClass);
            }
        }

        @Override
        public boolean commit(boolean dryRun) throws IOException, TerseException {
            boolean changed = false;
            for (Map.Entry<Path, Map<PluginType, Set<String>>> manifestSource : manifests.entrySet()) {
                Path jarPath = manifestSource.getKey();
                Map<PluginType, Set<String>> before = initial.get(jarPath);
                Map<PluginType, Set<String>> after = manifestSource.getValue();
                if (startSync(dryRun, jarPath, before, after)) {
                    rewriteJar(dryRun, jarPath, after);
                    changed = true;
                }
            }
            return changed;
        }
    }

    /**
     * The class hierarchy is similar to the single-jar because there can only be one manifest per type.
     * However, the path to that single manifest is accessed via the pluginLocation.
     */
    private class ClassHierarchyWorkspace extends SingleJarWorkspace {

        private ClassHierarchyWorkspace(PluginSource source) throws IOException {
            super(source);
        }

        @Override
        protected Map<PluginType, Set<String>> load(PluginSource source) throws IOException {
            return loadManifest(source.location().toUri().toURL());
        }

        protected boolean commit(boolean dryRun) throws IOException, TerseException {
            if (startSync(dryRun, location(), initial, manifests)) {
                rewriteClassHierarchyManifest(dryRun, location(), manifests);
                return true;
            }
            return false;
        }
    }

    private boolean startSync(boolean dryRun, Path syncLocation, Map<PluginType, Set<String>> before, Map<PluginType, Set<String>> after) {
        Objects.requireNonNull(syncLocation, "syncLocation must be non-null");
        Objects.requireNonNull(before, "before must be non-null");
        Objects.requireNonNull(after, "after must be non-null");
        if (before.equals(after)) {
            return false;
        }
        Set<String> added = new HashSet<>();
        after.values().forEach(added::addAll);
        before.values().forEach(added::removeAll);
        Set<String> removed = new HashSet<>();
        before.values().forEach(removed::addAll);
        after.values().forEach(removed::removeAll);
        out.printf("%sSync\t%s Add %s Remove %s%n", dryRun ? "Dry Run " : "", syncLocation, added.size(), removed.size());
        for (String add : added) {
            out.printf("\tAdd\t%s%n", add);
        }
        for (String rem : removed) {
            out.printf("\tRemove\t%s%n", rem);
        }
        return true;
    }

    /**
     * Rewrite a jar on disk to have manifests with the specified entries.
     * Will create the jar file if it does not exist and at least one manifest element is specified.
     * Will delete the jar file if it exists, and would otherwise remain empty.
     *
     * @param dryRun           True if the rewrite should be applied, false if it should be simulated.
     * @param jarPath          Path to a jar file for a plugin
     * @param manifestElements Map from plugin type to Class names of plugins which should appear in that manifest
     */
    private void rewriteJar(boolean dryRun, Path jarPath, Map<PluginType, Set<String>> manifestElements) throws IOException, TerseException {
        Objects.requireNonNull(jarPath, "jarPath must be non-null");
        Objects.requireNonNull(manifestElements, "manifestState must be non-null");
        Path writableJar = getWritablePath(dryRun, jarPath);
        if (nonEmpty(manifestElements) && !Files.exists(writableJar)) {
            log.debug("Create {}", jarPath);
            createJar(writableJar);
        }
        try (FileSystem jar = FileSystems.newFileSystem(
                new URI("jar", writableJar.toUri().toString(), ""),
                Map.of()
        )) {
            Path zipRoot = jar.getRootDirectories().iterator().next();
            // Set dryRun to false because this jar file is already a writable copy.
            rewriteClassHierarchyManifest(false, zipRoot, manifestElements);
        } catch (URISyntaxException e) {
            throw new IOException(e);
        }
        if (Files.exists(writableJar) && jarIsEmpty(writableJar)) {
            Files.delete(writableJar);
        }
    }

    private static boolean nonEmpty(Map<PluginType, Set<String>> manifestElements) {
        return !manifestElements.values().stream().allMatch(Collection::isEmpty);
    }

    private void createJar(Path path) throws IOException {
        Objects.requireNonNull(path, "path must be non-null");
        try (ZipOutputStream stream = new ZipOutputStream(Files.newOutputStream(
                path,
                StandardOpenOption.CREATE,
                StandardOpenOption.WRITE,
                StandardOpenOption.TRUNCATE_EXISTING
        ))) {
            stream.closeEntry();
        }
    }

    private boolean jarIsEmpty(Path path) throws IOException {
        Objects.requireNonNull(path, "path must be non-null");
        try (ZipInputStream stream = new ZipInputStream(Files.newInputStream(
                path,
                StandardOpenOption.READ
        ))) {
            return stream.getNextEntry() == null;
        }
    }

    /**
     * Rewrite multiple manifest files contained inside a class hierarchy.
     * Will create the files and parent directories if they do not exist and at least one element is specified.
     * Will delete the files and parent directories if they exist and would otherwise remain empty.
     *
     * @param dryRun           True if the rewrite should be applied, false if it should be simulated.
     * @param pluginLocation   Path to top-level of the class hierarchy
     * @param manifestElements Map from plugin type to Class names of plugins which should appear in that manifest
     */
    private void rewriteClassHierarchyManifest(boolean dryRun, Path pluginLocation, Map<PluginType, Set<String>> manifestElements) throws IOException, TerseException {
        Objects.requireNonNull(pluginLocation, "pluginLocation must be non-null");
        Objects.requireNonNull(manifestElements, "manifestState must be non-null");
        if (!Files.exists(pluginLocation)) {
            throw new TerseException(pluginLocation + " does not exist");
        }
        if (!Files.isWritable(pluginLocation)) {
            throw new TerseException(pluginLocation + " is not writable");
        }
        Path metaInfPath = pluginLocation.resolve("META-INF");
        Path servicesPath = metaInfPath.resolve("services");
        if (nonEmpty(manifestElements) && !Files.exists(servicesPath) && !dryRun) {
            Files.createDirectories(servicesPath);
        }
        for (Map.Entry<PluginType, Set<String>> manifest : manifestElements.entrySet()) {
            PluginType type = manifest.getKey();
            Set<String> elements = manifest.getValue();
            rewriteManifestFile(dryRun, servicesPath.resolve(type.superClass().getName()), elements);
        }
        deleteDirectoryIfEmpty(dryRun, servicesPath);
        deleteDirectoryIfEmpty(dryRun, metaInfPath);
    }

    private void deleteDirectoryIfEmpty(boolean dryRun, Path path) throws IOException, TerseException {
        if (!Files.exists(path)) {
            return;
        }
        if (!Files.isWritable(path)) {
            throw new TerseException(path + " is not writable");
        }
        try (Stream<Path> list = Files.list(path)) {
            if (list.findAny().isPresent()) {
                return;
            }
        }
        log.debug("Delete {}", path);
        if (!dryRun) {
            Files.delete(path);
        }
    }

    /**
     * Rewrite a single manifest file.
     * Will create the file if it does not exist and at least one element is specified.
     * Will delete the file if it exists and no elements are specified.
     *
     * @param dryRun              True if the rewrite should be applied, false if it should be simulated.
     * @param filePath            Path to file which should be rewritten.
     * @param elements            Class names of plugins which should appear in the manifest
     */
    private void rewriteManifestFile(boolean dryRun, Path filePath, Set<String> elements) throws IOException, TerseException {
        Objects.requireNonNull(filePath, "filePath must be non-null");
        Objects.requireNonNull(elements, "elements must be non-null");
        Path writableFile = getWritablePath(dryRun, filePath);
        if (elements.isEmpty()) {
            if (Files.exists(filePath)) {
                log.debug("Delete {}", filePath);
                if (!dryRun) {
                    Files.delete(writableFile);
                }
            }
        } else {
            if (!Files.exists(filePath)) {
                log.debug("Create {}", filePath);
            }
            log.debug("Write {} with content {}", filePath, elements);
            if (!dryRun) {
                try (OutputStream stream = new BufferedOutputStream(Files.newOutputStream(
                        writableFile,
                        StandardOpenOption.CREATE,
                        StandardOpenOption.WRITE,
                        StandardOpenOption.TRUNCATE_EXISTING
                ))) {
                    byte[] newline = System.lineSeparator().getBytes(StandardCharsets.UTF_8);
                    stream.write(MANIFEST_HEADER.getBytes(StandardCharsets.UTF_8));
                    stream.write(newline);
                    for (String element : elements) {
                        stream.write(element.getBytes(StandardCharsets.UTF_8));
                        stream.write(newline);
                    }
                }
            }
        }
    }

    /**
     * Get a path which is always writable
     * @param dryRun If true, substitute a temporary file instead of the real file on disk.
     * @param path Path which must be writable
     * @return Path which is writable, and may be different from the input path
     */
    private Path getWritablePath(boolean dryRun, Path path) throws IOException, TerseException {
        Objects.requireNonNull(path, "path must be non-null");
        for (Path parent = path; parent != null && !Files.isWritable(parent); parent = parent.getParent()) {
            if (Files.exists(parent) && !Files.isWritable(parent)) {
                throw new TerseException("Path " + path + " must be writable");
            }
        }
        if (dryRun) {
            if (!temporaryOverlayFiles.containsKey(path)) {
                Path fileName = path.getFileName();
                String suffix = fileName != null ? fileName.toString() : ".temp";
                Path temp = Files.createTempFile("connect-plugin-path-temporary-", suffix);
                if (Files.exists(path)) {
                    Files.copy(path, temp, StandardCopyOption.REPLACE_EXISTING);
                    temp.toFile().deleteOnExit();
                } else {
                    Files.delete(temp);
                }
                temporaryOverlayFiles.put(path, temp);
                return temp;
            }
            return temporaryOverlayFiles.get(path);
        }
        return path;
    }

    // Based on implementation from ServiceLoader.LazyClassPathLookupIterator from OpenJDK11
    private static Set<String> parse(URL u) {
        Set<String> names = new LinkedHashSet<>(); // preserve insertion order
        try {
            URLConnection uc = u.openConnection();
            uc.setUseCaches(false);
            try (InputStream in = uc.getInputStream();
                 BufferedReader r
                         = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))) {
                int lc = 1;
                while ((lc = parseLine(u, r, lc, names)) >= 0) {
                    // pass
                }
            }
        } catch (IOException x) {
            throw new RuntimeException("Error accessing configuration file", x);
        }
        return names;
    }

    // Based on implementation from ServiceLoader.LazyClassPathLookupIterator from OpenJDK11
    private static int parseLine(URL u, BufferedReader r, int lc, Set<String> names) throws IOException {
        String ln = r.readLine();
        if (ln == null) {
            return -1;
        }
        int ci = ln.indexOf('#');
        if (ci >= 0) ln = ln.substring(0, ci);
        ln = ln.trim();
        int n = ln.length();
        if (n != 0) {
            if ((ln.indexOf(' ') >= 0) || (ln.indexOf('\t') >= 0))
                throw new IOException("Illegal configuration-file syntax in " + u);
            int cp = ln.codePointAt(0);
            if (!Character.isJavaIdentifierStart(cp))
                throw new IOException("Illegal provider-class name: " + ln + " in " + u);
            int start = Character.charCount(cp);
            for (int i = start; i < n; i += Character.charCount(cp)) {
                cp = ln.codePointAt(i);
                if (!Character.isJavaIdentifierPart(cp) && (cp != '.'))
                    throw new IOException("Illegal provider-class name: " + ln + " in " + u);
            }
            names.add(ln);
        }
        return lc + 1;
    }
}
