LibJF/libjf-config-compiler-plugin/src/main/java/io/gitlab/jfronny/libjf/config/plugin/StreamAction.java

package io.gitlab.jfronny.libjf.config.plugin;

import io.gitlab.jfronny.gson.stream.JsonReader;
import io.gitlab.jfronny.gson.stream.JsonWriter;
import io.gitlab.jfronny.libjf.config.plugin.asm.ConfigInjectClassTransformer;
import io.gitlab.jfronny.libjf.config.plugin.asm.NotAConfigClassException;
import io.gitlab.jfronny.libjf.config.plugin.fmj.FabricModJsonTransformer;
import io.gitlab.jfronny.libjf.config.plugin.util.*;
import org.apache.tools.zip.*;
import org.gradle.api.GradleException;
import org.gradle.api.file.FileCopyDetails;
import org.objectweb.asm.*;

import java.io.*;
import java.util.GregorianCalendar;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

public class StreamAction {
    private static final long CONSTANT_TIME_FOR_ZIP_ENTRIES = (new GregorianCalendar(1980, 1, 1, 0, 0, 0)).getTimeInMillis();
    private final ZipOutputStream zipOutStr;
    private final File zipFile;
    private final boolean preserveFileTimestamps;
    private final String modId;
    private final Set<Type> knownConfigClasses;

    public StreamAction(ZipOutputStream zipOutStr, File zipFile, boolean preserveFileTimestamps, String modId, Set<Type> knownConfigClasses) {
        this.zipOutStr = zipOutStr;
        this.zipFile = zipFile;
        this.preserveFileTimestamps = preserveFileTimestamps;
        this.modId = modId;
        this.knownConfigClasses = knownConfigClasses;
    }

    public boolean processFile(FileCopyDetails details) {
        try {
            if (details.isDirectory()) {
                visitDirectory(details);
                return false;
            } else {
                return visitFile(details);
            }
        } catch (Throwable e) {
            throw new GradleException("Could not add " + details + " to ZIP " + zipFile, e);
        }
    }

    private void visitDirectory(FileCopyDetails details) throws IOException {
        String path = details.getRelativePath().getPathString() + "/";
        ZipEntry archiveEntry = new ZipEntry(path);
        archiveEntry.setTime(getArchiveTimeFor(details.getLastModified()));
        archiveEntry.setUnixMode(UnixStat.DIR_FLAG | details.getMode());
        zipOutStr.putNextEntry(archiveEntry);
        zipOutStr.closeEntry();
    }

    private boolean visitFile(FileCopyDetails details) throws IOException {
        if (details.getPath().endsWith(".jar")) {
            return processArchive(details);
        } else if (details.getPath().endsWith(".class")) {
            processClass(details);
        } else if (details.getPath().equals("fabric.mod.json")) {
            processFMJ(details);
            return true;
        } else {
            ZipEntry archiveEntry = new ZipEntry(details.getRelativePath().getPathString());
            archiveEntry.setTime(getArchiveTimeFor(details.getLastModified()));
            archiveEntry.setUnixMode(UnixStat.FILE_FLAG | details.getMode());
            zipOutStr.putNextEntry(archiveEntry);
            details.copyTo(zipOutStr);
            zipOutStr.closeEntry();
        }
        return false;
    }

    private boolean processArchive(FileCopyDetails details) throws IOException {
        try (ZipFile archive = new ZipFile(details.getFile())) {
            AtomicBoolean fmjFound = new AtomicBoolean(false);
            EnumerationSpliterator.stream(archive.getEntries())
                    .map(RelativeArchivePath::new)
                    .map(ArchiveFileTreeElement::new)
                    .filter(it -> it.getRelativePath().isFile())
                    .forEach(archiveElement -> {
                        fmjFound.compareAndSet(false, visitArchiveFile(archiveElement, archive));
                    });
            return fmjFound.get();
        }
    }

    private void visitArchiveDirectory(RelativeArchivePath archiveDir) throws IOException {
        zipOutStr.putNextEntry(archiveDir.entry);
        zipOutStr.closeEntry();
    }

    private boolean visitArchiveFile(ArchiveFileTreeElement archiveFile, ZipFile archive) {
        RelativeArchivePath archiveFilePath = archiveFile.getRelativePath();
        try {
            if (archiveFile.isClassFile()) {
                processClass(archiveFilePath, archive);
            } else if (archiveFilePath.getPathString().equals("fabric.mod.json")) {
                processFMJ(archiveFilePath, archive);
                return true;
            } else {
                copyArchiveEntry(archiveFilePath, archive);
            }
            return false;
        } catch (Throwable e) {
            throw new GradleException("Could not read archive entry " + archiveFilePath.getPathString(), e);
        }
    }

    private void copyArchiveEntry(RelativeArchivePath archiveFile, ZipFile archive) throws IOException {
        ZipEntry entry = new ZipEntry(archiveFile.entry.getName());
        entry.setTime(getArchiveTimeFor(archiveFile.entry.getTime()));
        RelativeArchivePath path = new RelativeArchivePath(entry);
        addParentDirectories(path);
        zipOutStr.putNextEntry(path.entry);
        try (InputStream is = archive.getInputStream(archiveFile.entry)) {
            byte[] buffer = new byte[1024];
            int n;
            while (-1 != (n = is.read(buffer))) {
                zipOutStr.write(buffer, 0, n);
            }
        }
        zipOutStr.closeEntry();
    }

    private void addParentDirectories(RelativeArchivePath file) throws IOException {
        if (file != null) {
            addParentDirectories(file.getParent());
            if (!file.isFile()) {
                visitArchiveDirectory(file);
            }
        }
    }

    private void processClass(RelativeArchivePath file, ZipFile archive) throws IOException {
        ZipEntry zipEntry = new ZipEntry(file.getPathString());
        addParentDirectories(new RelativeArchivePath(zipEntry));
        processClass(archive.getInputStream(file.entry).readAllBytes(), file.getPathString(), file.entry.getTime());
    }

    private void processClass(FileCopyDetails details) throws IOException {
        try (InputStream is = new FileInputStream(details.getFile())) {
            processClass(is.readAllBytes(), details.getPath(), details.getLastModified());
        }
    }

    private void processClass(byte[] klazz, String path, long lastModified) throws IOException {
        try {
            final ClassReader reader = new ClassReader(klazz);
            final ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES);
            final ConfigInjectClassTransformer transformer = new ConfigInjectClassTransformer(modId, writer, knownConfigClasses);
            reader.accept(transformer, ClassReader.EXPAND_FRAMES);
            klazz = writer.toByteArray();
            System.out.println("Injected config registration into " + path);
        } catch (NotAConfigClassException notAConfigClass) {
            // Use original bytes
        }
        ZipEntry archiveEntry = new ZipEntry(path);
        archiveEntry.setTime(getArchiveTimeFor(lastModified));
        zipOutStr.putNextEntry(archiveEntry);
        zipOutStr.write(klazz);
        zipOutStr.closeEntry();
    }

    private void processFMJ(RelativeArchivePath file, ZipFile archive) throws IOException {
        ZipEntry zipEntry = new ZipEntry(file.getPathString());
        addParentDirectories(new RelativeArchivePath(zipEntry));
        processFMJ(archive.getInputStream(file.entry), file.getPathString(), file.entry.getTime());
    }

    private void processFMJ(FileCopyDetails details) throws IOException {
        try (InputStream is = new FileInputStream(details.getFile())) {
            processFMJ(is, details.getPath(), details.getLastModified());
        }
    }

    private void processFMJ(InputStream fmj, String path, long lastModified) throws IOException {
        ZipEntry archiveEntry = new ZipEntry(path);
        archiveEntry.setTime(getArchiveTimeFor(lastModified));
        zipOutStr.putNextEntry(archiveEntry);
        // Leave this closeable open, as everything else will break the writer

        try (JsonReader reader = new JsonReader(new InputStreamReader(fmj));
             JsonWriter writer = new JsonWriter(new OutputStreamWriter(new DelegatingUncloseableOutputStream(zipOutStr)))) {
            FabricModJsonTransformer.transform(reader, writer, knownConfigClasses);
        }
        zipOutStr.closeEntry();
    }

    private long getArchiveTimeFor(long timestamp) {
        return preserveFileTimestamps ? timestamp : CONSTANT_TIME_FOR_ZIP_ENTRIES;
    }
}