path: root/src/main/java/org/zwobble/hobgoblin/compiler/output/CodeWriter.java

package org.zwobble.hobgoblin.compiler.output;

import java.io.FileWriter;
import java.io.IOException;
import java.io.StringWriter;
import java.io.Writer;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.util.*;

public class CodeWriter implements AutoCloseable {
    public interface CodeWriterFile {
        String readString() throws IOException;
        Writer openWriter() throws IOException;
    }

    public static class InMemoryCodeWriterFile implements CodeWriterFile {
        private final Optional<String> originalContents;
        private final StringWriter stringWriter;
        private boolean writerOpened;

        public InMemoryCodeWriterFile(
            Optional<String> originalContents,
            StringWriter stringWriter
        ) {
            this.originalContents = originalContents;
            this.stringWriter = stringWriter;
            this.writerOpened = false;
        }

        @Override
        public String readString() throws IOException {
            if (originalContents.isPresent()) {
                return originalContents.get();
            } else {
                throw new NoSuchFileException("");
            }
        }

        @Override
        public Writer openWriter() throws IOException {
            this.writerOpened = true;
            return stringWriter;
        }

        public boolean isWriterOpened() {
            return writerOpened;
        }
    }

    public record CustomAreas(
        Map<String, List<String>> existing,
        String customAreaStart,
        String customAreaEnd
    ) {
        public static CustomAreas read(
            String existingString,
            String customAreaStart,
            String customAreaEnd
        ) {
            var existing = new HashMap<String, List<String>>();

            Iterable<String> lines = () -> existingString.lines().iterator();
            Optional<String> name = Optional.empty();
            var customAreaLines = new ArrayList<String>();
            for (var line : lines) {
                var trimmedLine = line.trim();
                if (trimmedLine.startsWith(customAreaStart)) {
                    name = Optional.of(trimmedLine.substring(customAreaStart.length()));
                } else if (name.isPresent() && trimmedLine.startsWith(customAreaEnd + name.get())) {
                    existing.put(name.get(), customAreaLines);
                    customAreaLines = new ArrayList<>();
                    name = Optional.empty();
                } else if (name.isPresent()) {
                    customAreaLines.add(line);
                }
            }

            return new CustomAreas(existing, customAreaStart, customAreaEnd);
        }

        public static CustomAreas empty(
            String customAreaStart,
            String customAreaEnd
        ) {
            return new CustomAreas(Map.of(), customAreaStart, customAreaEnd);
        }

        public Iterable<String> lines(String name) {
            var lines = this.existing.get(name);
            if (lines == null) {
                return List.of();
            } else {
                return lines;
            }
        }
    }

    public interface Write {
        void write(CodeWriter writer) throws IOException;
    }

    public static void writeTo(
        Path path,
        String lineCommentStart,
        Write write
    ) throws IOException {
        writeTo(
            new CodeWriterFile() {
                @Override
                public String readString() throws IOException {
                    return Files.readString(path);
                }

                @Override
                public Writer openWriter() throws IOException {
                    Files.createDirectories(path.getParent());
                    return new FileWriter(path.toFile());
                }
            },
            lineCommentStart,
            write
        );
    }

    public static void writeTo(
        CodeWriterFile file,
        String lineCommentStart,
        Write write
    ) throws IOException {
        var originalContents = readFile(file);

        var header = lineCommentStart + " Generated by hobgoblin.\n\n";
        if (originalContents.isPresent() && !originalContents.get().startsWith(header)) {
            return;
        }

        var customAreaStart = lineCommentStart + " Custom area start: ";
        var customAreaEnd = lineCommentStart + " Custom area end: ";

        var customAreas = readCustomAreas(file, customAreaStart, customAreaEnd);
        try (var writer = file.openWriter()) {
            // TODO: warn if custom areas are unused?
            write.write(new CodeWriter(writer, header, customAreas));
        }
    }

    private static CustomAreas readCustomAreas(
        CodeWriterFile file,
        String customAreaStart,
        String customAreaEnd
    ) throws IOException {
        try {
            var contents = file.readString();
            return CustomAreas.read(contents, customAreaStart, customAreaEnd);
        } catch (NoSuchFileException exception) {
            return CustomAreas.empty(customAreaStart, customAreaEnd);
        }
    }

    private static Optional<String> readFile(CodeWriterFile file) throws IOException {
        try {
            return Optional.of(file.readString());
        } catch (NoSuchFileException exception) {
            return Optional.empty();
        }
    }

    private final Writer writer;
    private final String header;
    private final CustomAreas customAreas;
    private int indent;
    private boolean pendingIndent;

    public CodeWriter(Writer writer, String header, CustomAreas customAreas) {
        this.writer = writer;
        this.header = header;
        this.customAreas = customAreas;
        this.indent = 0;
        this.pendingIndent = false;
    }

    public void writeHeader() throws IOException {
        writer.write(header);
    }

    public void write(String value) throws IOException {
        if (this.pendingIndent) {
            for (var i = 0; i < this.indent; i++) {
                this.writer.write("    ");
            }
            this.pendingIndent = false;
        }

        this.writer.write(value);
    }

    public void newLine() throws IOException {
        this.writer.write("\n");
        this.pendingIndent = true;
    }

    public void indent() {
        this.indent += 1;
    }

    public void dedent() {
        this.indent -= 1;
    }

    public void writeCustomArea(String name) throws IOException {
        this.write(this.customAreas.customAreaStart);
        this.write(name);

        for (var line : this.customAreas.lines(name)) {
            this.writer.write("\n");
            this.write(line);
        }

        this.newLine();
        this.write(this.customAreas.customAreaEnd);
        this.write(name);
        this.newLine();
    }

    @Override
    public void close() throws IOException {
        this.writer.close();
    }

    public static <T> void writeWithSeparator(
        Iterable<T> iterable,
        WriteElement<T> writeElement,
        WriteSeparator writeSeparator
    ) throws IOException {
        var isFirst = true;
        for (var element : iterable) {
            if (!isFirst) {
                writeSeparator.write();
            }

            writeElement.write(element);

            if (isFirst) {
                isFirst = false;
            }
        }
    }

    public interface WriteElement<T> {
        void write(T element) throws IOException;
    }

    public interface WriteSeparator {
        void write() throws IOException;
    }
}