package org.zwobble.hobgoblin.compiler.output;

import org.zwobble.hobgoblin.compiler.output.generators.Stripper;

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 static class CodeWriterConfig {
        private final String header;
        private final String customAreaStart;
        private final String customAreaEnd;

        public CodeWriterConfig(String lineCommentStart) {
            this.header = lineCommentStart + " Generated by hobgoblin.\n\n";
            this.customAreaStart = lineCommentStart + " Custom area start: ";
            this.customAreaEnd = lineCommentStart + " Custom area end: ";
        }

        public String header() {
            return this.header;
        }

        public String customAreaStart() {
            return customAreaStart;
        }

        public String customAreaEnd() {
            return customAreaEnd;
        }
    }

    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(
        SequencedMap<String, List<String>> existing,
        CodeWriterConfig config
    ) {
        public static CustomAreas read(
            String existingString,
            CodeWriterConfig config
        ) {
            var existing = new LinkedHashMap<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(config.customAreaStart)) {
                    name = Optional.of(trimmedLine.substring(config.customAreaStart.length()));
                } else if (name.isPresent() && trimmedLine.startsWith(config.customAreaEnd + name.get())) {
                    if (!customAreaLines.isEmpty()) {
                        existing.put(name.get(), customAreaLines);
                    }
                    customAreaLines = new ArrayList<>();
                    name = Optional.empty();
                } else if (name.isPresent()) {
                    customAreaLines.add(line);
                }
            }

            return new CustomAreas(existing, config);
        }

        public static CustomAreas empty(CodeWriterConfig config) {
            return new CustomAreas(new LinkedHashMap<>(), config);
        }

        public String customAreaStart() {
            return this.config.customAreaStart;
        }

        public String customAreaEnd() {
            return this.config.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,
        CodeWriterConfig config,
        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());
                }
            },
            config,
            write
        );
    }

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

        if (originalContents.isPresent() && !isGeneratedFile(originalContents.get(), config)) {
            return;
        }

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

    private static boolean isGeneratedFile(String contents, CodeWriterConfig config) {
        return contents.startsWith(config.header);
    }

    public static Stripper stripper(
        Path outputPath,
        CodeWriterConfig config
    ) {
        return new Stripper(outputPath, config);
    }

    public static void strip(Path outputPath, CodeWriterConfig config) throws IOException {
        try (var filesStream = Files.walk(outputPath)) {
            var filesIterator = filesStream.iterator();
            while (filesIterator.hasNext()) {
                var path = filesIterator.next();
                if (path.toFile().isFile()) {
                    // TODO: filter by extension?
                    // TODO: don't read entire string if not necessary.
                    var contents = Files.readString(path);
                    switch (stripFileContents(contents, config)) {
                        case StripResult.Delete _ -> {
                            Files.delete(path);
                        }
                        case StripResult.NewContent newContent -> {
                            Files.writeString(path, newContent.newContents);
                        }
                        case StripResult.Unchanged _ -> {
                        }
                    }
                }
            }
        }
    }

    sealed interface StripResult {
        record Unchanged() implements StripResult {}
        record Delete() implements StripResult {}
        record NewContent(String newContents) implements StripResult {}
    }

    static StripResult stripFileContents(String contents, CodeWriterConfig config) {
        if (!isGeneratedFile(contents, config)) {
            return new StripResult.Unchanged();
        }

        var stringWriter = new StringWriter();
        var customAreas = CustomAreas.read(contents, config);
        if (customAreas.existing.isEmpty()) {
            return new StripResult.Delete();
        }

        var writer = new CodeWriter(stringWriter, config.header, customAreas);
        try {
            writer.writeHeader();
            for (var customAreaName : customAreas.existing.keySet()) {
                writer.writeCustomArea(customAreaName);
            }
            return new StripResult.NewContent(stringWriter.toString());
        } catch (IOException exception) {
            // In theory, writing to a StringWriter will always succeed, so
            // we don't need to handle this especially well.
            throw new RuntimeException(exception);
        }
    }

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

    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;
    }
}