package org.shedlang.compiler.backends.cranelift;

import org.pcollections.PVector;
import org.shedlang.compiler.backends.Backend;
import org.shedlang.compiler.backends.PatioBackends;
import org.shedlang.compiler.backends.llvm.LlvmBackend;
import org.shedlang.compiler.backends.paths.PartialTargetPath;
import org.shedlang.compiler.backends.paths.TargetDir;
import org.shedlang.compiler.backends.targets.NativeTarget;
import org.shedlang.compiler.config.ProfileConfig;
import org.shedlang.compiler.config.ProjectConfig;
import org.shedlang.compiler.filesystem.ShedCompilerFiles;
import org.shedlang.compiler.patio.PatioCompilationUnit;
import org.shedlang.compiler.patio.serialization.PatioToTransient0;
import org.shedlang.compiler.runtime.ShedRuntime;
import org.shedlang.compiler.shid.ShidNamespace;
import org.shedlang.compiler.typechecker.TypeInfos;
import org.shedlang.compiler.types.NamespaceName;
import org.shedlang.compiler.util.P;
import org.shedlang.compiler.util.processes.ProcessError;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

import static org.shedlang.compiler.util.processes.ProcessRunner.runInheritIo;

public class CraneliftBackend implements Backend {
    public static final CraneliftBackend AMD64 = forTarget(NativeTarget.AMD64);

    public static CraneliftBackend forTarget(NativeTarget target) {
        return new CraneliftBackend(target);
    }

    private final NativeTarget target;

    private CraneliftBackend(NativeTarget target) {
        this.target = target;
    }

    public Path compileCompilationUnit(
        PatioCompilationUnit patio,
        PartialTargetPath partialTargetPath
    ) throws IOException, InterruptedException {
        return compileCompilationUnit(patio, Optional.empty(), partialTargetPath);
    }

    private Path compileMain(
        PatioCompilationUnit patio,
        String mainFunctionName,
        PartialTargetPath partialTargetPath
    ) throws IOException, InterruptedException {
        return compileCompilationUnit(patio, Optional.of(mainFunctionName), partialTargetPath);
    }

    private Path compileCompilationUnit(
        PatioCompilationUnit patio,
        Optional<String> mainFunctionName,
        PartialTargetPath partialTargetPath
    ) throws IOException, InterruptedException {
        var patioPath = partialTargetPath.withExtension(".patio");
        Files.createDirectories(patioPath.getParent());
        try (var outputStream = new FileOutputStream(patioPath.toFile())) {
            PatioToTransient0.serializeCompilationUnit(patio, outputStream);
        }

        var patioToCraneliftArgs = new ArrayList<String>();
        patioToCraneliftArgs.add(
            ShedCompilerFiles.getRoot()
                .resolve("patio-to-cranelift/target/debug/patio-to-cranelift")
                .toString()
        );
        patioToCraneliftArgs.add(patioPath.toString());
        patioToCraneliftArgs.add("--target");
        patioToCraneliftArgs.add(this.target.targetTriple());

        if (mainFunctionName.isPresent()) {
            patioToCraneliftArgs.add("--main");
            patioToCraneliftArgs.add(mainFunctionName.get());
        }

        runInheritIo(patioToCraneliftArgs);

        var objectFilePath = partialTargetPath.withExtension(".o");
        return objectFilePath;
    }

    public void link(
        List<Path> inputs,
        Path outputPath,
        ShedRuntime shedRuntime
    ) throws ProcessError, InterruptedException {
        // TODO: extract linking if it's not specific to LLVM.
        LlvmBackend.forTarget(this.target).link(inputs, outputPath, shedRuntime);
    }

    @Override
    public String targetName() {
        return "cranelift-" + target.targetName();
    }

    @Override
    public void compile(
        ProjectConfig config,
        ProfileConfig profileConfig,
        List<ShidNamespace> namespaces,
        TypeInfos typeInfos,
        TargetDir targetDir
    ) throws IOException, InterruptedException {
        PatioBackends.compile(
            config,
            profileConfig,
            namespaces,
            typeInfos,
            new PatioToBinary(
                targetDir,
                P.vector()
            )
        );
    }

    private class PatioToBinary implements PatioBackends.PatioToTarget {
        private final TargetDir targetDir;
        private final PVector<Path> objectFilePaths;

        public PatioToBinary(
            TargetDir targetDir,
            PVector<Path> objectFilePaths
        ) {
            this.targetDir = targetDir;
            this.objectFilePaths = objectFilePaths;
        }

        @Override
        public PatioBackends.PatioToTarget namespace(
            NamespaceName namespaceName,
            PatioCompilationUnit patioCompilationUnit
        ) throws IOException, InterruptedException {
            var objectPath = compileCompilationUnit(
                patioCompilationUnit,
                this.targetDir.namespace(namespaceName)
            );

            return new PatioToBinary(
                this.targetDir,
                this.objectFilePaths.plus(objectPath)
            );
        }

        @Override
        public void bin(
            String binName,
            PatioCompilationUnit patioCompilationUnit,
            String mainFunctionName,
            ProfileConfig profileConfig
        ) throws IOException, InterruptedException {
            var linkInputs = new ArrayList<>(objectFilePaths);

            var mainObjectPath = compileMain(
                patioCompilationUnit,
                mainFunctionName,
                this.targetDir.bin(binName)
            );
            linkInputs.add(mainObjectPath);

            var shedRuntime = ShedRuntime.runtime(
                target.targetName(),
                profileConfig.gc(),
                targetDir
            );

            link(
                linkInputs,
                targetDir.topLevel(binName),
                shedRuntime
            );
        }
    }
}