package org.shedlang.compiler.backends.llvm;

import org.pcollections.PVector;
import org.shedlang.compiler.backends.Backend;
import org.shedlang.compiler.backends.PatioBackends;
import org.shedlang.compiler.backends.llvm.codegen.PatioToLlvm;
import org.shedlang.compiler.backends.llvm.codegen.PatioToLlvmGlobalContext;
import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDefinition;
import org.shedlang.compiler.backends.llvm.ir.ast.LlvmModule;
import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmCall;
import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmReturn;
import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmTrunc;
import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
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.patio.PatioCompilationUnit;
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.IOException;
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 LlvmBackend implements Backend {
    public static final LlvmBackend AMD64 = new LlvmBackend(NativeTarget.AMD64);
    public static final LlvmBackend I686 = new LlvmBackend(NativeTarget.I686);

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

    private final NativeTarget target;

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

    @Override
    public String targetName() {
        return this.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(),
                PatioToLlvmGlobalContext.empty(this.target)
            )
        );
    }

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

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

        @Override
        public PatioBackends.PatioToTarget namespace(
            NamespaceName namespaceName,
            PatioCompilationUnit patioCompilationUnit
        ) throws IOException, InterruptedException {
            var llvmModuleResult = PatioToLlvm.generateCompilationUnit(
                patioCompilationUnit,
                this.patioToLlvmGlobalContext
            );
            var llvmModule = llvmModuleResult.value();
            var patioToLlvmGlobalContext = llvmModuleResult.context();

            var objectPath = LlvmCompiler.compileLlvmModule(
                llvmModule,
                targetDir.namespace(namespaceName)
            );

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

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

            var mainModule = generateMain(
                patioCompilationUnit,
                mainFunctionName,
                patioToLlvmGlobalContext
            );
            var mainObjectPath = LlvmCompiler.compileLlvmModule(mainModule, targetDir.bin(binName));
            linkInputs.add(mainObjectPath);

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

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

    public void link(
        List<Path> inputs,
        Path outputPath,
        ShedRuntime shedRuntime
    ) throws ProcessError, InterruptedException {
        var clangArgs = new ArrayList<String>();
        clangArgs.add("clang");

        clangArgs.add("-o");
        clangArgs.add(outputPath.toString());

        for (var input : inputs) {
            clangArgs.add(input.toString());
        }

        clangArgs.add(shedRuntime.dir().resolve("libpatio.a").toString());
        clangArgs.add(shedRuntime.dir().resolve("libshed.a").toString());

        clangArgs.addAll(this.target.clangFlags());

        runInheritIo(clangArgs);
    }

    private static LlvmModule generateMain(
        PatioCompilationUnit patioCompilationUnit,
        String mainFunctionName,
        PatioToLlvmGlobalContext patioToLlvmGlobalContext
    ) {
        var llvmModuleResult = PatioToLlvm.generateCompilationUnit(
            patioCompilationUnit,
            patioToLlvmGlobalContext
        );
        var llvmModule = llvmModuleResult.value();

        llvmModule = llvmModule.addTopLevelEntity(new LlvmFunctionDefinition(
            "main",
            LlvmTypes.INT_32,
            List.of(),
            List.of(
                new LlvmBasicBlock(
                    Optional.empty(),
                    List.of(
                        new LlvmCall(
                            new LlvmVariableLocal("exitCode"),
                            LlvmTypes.INT_64,
                            new LlvmVariableGlobal(mainFunctionName),
                            List.of()
                        ),
                        new LlvmTrunc(
                            new LlvmVariableLocal("exitCode32"),
                            LlvmTypes.INT_64,
                            new LlvmVariableLocal("exitCode"),
                            LlvmTypes.INT_32
                        )
                    ),
                    new LlvmReturn(LlvmTypes.INT_32, new LlvmVariableLocal("exitCode32"))
                )
            )
        ));

        return llvmModule;
    }
}