package org.shedlang.compiler;

import org.shedlang.compiler.ast.untyped.UntypedNamespaceNode;
import org.shedlang.compiler.backends.Backend;
import org.shedlang.compiler.config.ConfigChecker;
import org.shedlang.compiler.config.ProjectConfig;
import org.shedlang.compiler.parser.Parser;
import org.shedlang.compiler.sources.Sources;
import org.shedlang.compiler.shid.ShidCodeGen;
import org.shedlang.compiler.shid.ShidNamespace;
import org.shedlang.compiler.typechecker.TypeChecker;
import org.shedlang.compiler.typechecker.TypeCheckerContext;
import org.shedlang.compiler.types.*;
import org.shedlang.compiler.util.P;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import static org.shedlang.compiler.parser.Tokeniser.tokenise;

public class ShedCompiler {
    public void compile(Path projectRoot, Backend backend) throws IOException, InterruptedException {
        var config = ProjectConfig.read(projectRoot);

        var sourceRoot = projectRoot.resolve("src");

        var sourceFilePaths = collectSourceFilePaths(sourceRoot);

        var untypedNamespaceNodes = parseSource(sourceRoot, sourceFilePaths);

        var namespaces = new ArrayList<ShidNamespace>();

        var typeCheckerContext = TypeCheckerContext.EMPTY
            .addNamespaceTypeInfo(
                ShedBuiltins.NAMESPACE_NAME,
                new NamespaceTypeInfo(FieldsInfo.of(
                    FieldsInfo.freeField(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
                ))
            )
            .enterGlobal(P.map(
                Map.entry("Bool", Types.metaType(Types.BOOL)),
                Map.entry("Int", Types.metaType(Types.INT)),
                Map.entry("Io", Types.effectType(Effects.IO)),
                Map.entry("String", Types.metaType(Types.STRING)),
                Map.entry("Unit", Types.metaType(Types.UNIT)),

                Map.entry(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
            ));

        for (var untypedNamespaceNode : orderSourceFiles(untypedNamespaceNodes)) {
            var typeCheckNamespaceResult = TypeChecker.typeCheckNamespace(untypedNamespaceNode, typeCheckerContext);
            typeCheckerContext = typeCheckNamespaceResult.context();

            var shidNamespace = ShidCodeGen.compileNamespace(typeCheckNamespaceResult.value());
            namespaces.add(shidNamespace);
        }

        var typeInfos = typeCheckerContext.toTypeInfos();

        ConfigChecker.configCheck(config, typeInfos);

        backend.compile(
            config,
            namespaces,
            typeInfos,
            config.targetDir().resolve(backend.name())
        );
    }

    private List<Path> collectSourceFilePaths(Path sourceRoot) throws IOException {
        var paths = new ArrayList<Path>();
        Files.walkFileTree(sourceRoot, new SimpleFileVisitor<Path>() {
            @Override
            public FileVisitResult visitFile(
                Path file,
                BasicFileAttributes attrs
            ) {
                if (file.getFileName().toString().endsWith(".shed")) {
                    paths.add(file);
                }
                return FileVisitResult.CONTINUE;
            }

            @Override
            public FileVisitResult postVisitDirectory(
                Path dir,
                IOException exc
            ) {
                return FileVisitResult.CONTINUE;
            }
        });
        paths.sort(Comparator.naturalOrder());
        return paths;
    }

    private List<UntypedNamespaceNode> parseSource(
        Path sourceRoot,
        List<Path> sourceFilePaths
    ) throws IOException {
        var untypedNamespaceNodes = new ArrayList<UntypedNamespaceNode>();

        for (var sourceFilePath : sourceFilePaths) {
            var namespaceNameParts = StreamSupport.stream(sourceRoot.relativize(sourceFilePath).spliterator(), false)
                .map(element -> {
                    var string = element.toString();
                    return string.substring(0, string.lastIndexOf('.'));
                })
                .toList();
            var namespaceName = new NamespaceName(namespaceNameParts);

            var sourceContents = Files.readString(sourceFilePath, StandardCharsets.UTF_8);
            var source = Sources.file(sourceFilePath.toString(), sourceContents);

            var tokens = tokenise(source);

            var untypedNode = Parser.parseNamespace(tokens, namespaceName);
            untypedNamespaceNodes.add(untypedNode);
        }

        return untypedNamespaceNodes;
    }

    private List<UntypedNamespaceNode> orderSourceFiles(
        List<UntypedNamespaceNode> untypedNamespaceNodes
    ) {
        var nodesByName = untypedNamespaceNodes.stream()
            .collect(Collectors.toMap(x -> x.namespaceName(), x -> x));

        var seen = new HashSet<NamespaceName>();
        var ordered = new ArrayList<UntypedNamespaceNode>();

        // TODO: test this!
        // TODO: handle infinite loops
        // TODO: handle missing namespaces

        var adder = new Object() {
            private void add(NamespaceName name) {
                if (!seen.contains(name)) {
                    var untypedNamespaceNode = nodesByName.get(name);
                    for (var import_ : untypedNamespaceNode.imports()) {
                        add(import_.namespaceName());
                    }
                    seen.add(name);
                    ordered.add(untypedNamespaceNode);
                }
            }
        };

        for (var untypedNamespaceNode : untypedNamespaceNodes) {
            adder.add(untypedNamespaceNode.namespaceName());
        }

        return ordered;
    }
}