package org.shedlang.compiler.backends.wasm;

import org.pcollections.PVector;
import org.shedlang.compiler.backends.Backend;
import org.shedlang.compiler.backends.PatioBackends;
import org.shedlang.compiler.backends.paths.PartialTargetPath;
import org.shedlang.compiler.backends.paths.TargetDir;
import org.shedlang.compiler.backends.targets.WasmTarget;
import org.shedlang.compiler.backends.wasm.structure.WasmBinaryFormat;
import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmCall;
import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmDrop;
import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmExpression;
import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFuncName;
import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFunction;
import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
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.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;

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

public class WasmBackend implements Backend {
    public static final WasmBackend WASM32 = new WasmBackend(WasmTarget.WASM32);
    public static final WasmBackend WASM64 = new WasmBackend(WasmTarget.WASM64);
    public static final String EXECUTABLE_EXTENSION = ".wasm";

    public static WasmBackend forTarget(WasmTarget target) {
        return new WasmBackend(target);
    }

    private final WasmTarget target;

    private WasmBackend(WasmTarget 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 PatioToWasmBinary(targetDir, P.vector())
        );
    }

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

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

        @Override
        public PatioBackends.PatioToTarget namespace(
            NamespaceName namespaceName,
            PatioCompilationUnit patioCompilationUnit
        ) throws IOException {
            var wasmModule = PatioToWasm.generateCompilationUnit(
                patioCompilationUnit,
                target
            );

            var objectPath = writeObjectFile(
                wasmModule,
                this.targetDir.namespace(namespaceName)
            );

            return new PatioToWasmBinary(
                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 mainModule = generateMain(patioCompilationUnit, mainFunctionName);
            var mainObjectPath = writeObjectFile(mainModule, targetDir.bin(binName));
            linkInputs.add(mainObjectPath);

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

            var outputBinPath = targetDir.topLevel(binName + EXECUTABLE_EXTENSION);
            link(
                linkInputs,
                outputBinPath,
                shedRuntime
            );
            validate(outputBinPath);
        }
    }

    public void link(
        List<Path> inputs,
        Path outputPath,
        ShedRuntime shedRuntime
    ) throws ProcessError, InterruptedException {
        var ldArgs = new ArrayList<String>();
        ldArgs.add("wasm-ld-23");

        ldArgs.addAll(this.target.linkerFlags());

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

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

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

        runInheritIo(ldArgs);
    }

    private WasmModule generateMain(
        PatioCompilationUnit patioCompilationUnit,
        String mainFunctionName
    ) {
        var wasmModule = PatioToWasm.generateCompilationUnit(
            patioCompilationUnit,
            this.target
        );

        wasmModule = generateNativeRuntime(wasmModule);

        wasmModule = wasmModule.addFunction(new WasmFunction(
            "_start",
            List.of(),
            List.of(),
            List.of(),
            new WasmExpression(List.of(
                new WasmCall(new WasmFuncName(mainFunctionName)),
                new WasmDrop()
            ))
        ));

        return wasmModule;
    }

    /**
     * Generates the parts of the runtime that are written in native WASM.
     */
    public static WasmModule generateNativeRuntime(WasmModule wasmModule) {
        return wasmModule;
    }

    private static Path writeObjectFile(
        WasmModule module,
        PartialTargetPath targetPath
    ) throws IOException {
        var objectPath = targetPath.withExtension(".o");
        Files.createDirectories(objectPath.getParent());
        try (var outputStream = new FileOutputStream(objectPath.toFile())) {
            WasmBinaryFormat.serializeModule(module, outputStream);
        }
        return objectPath;
    }

    private void validate(Path path) throws IOException, InterruptedException {
        var command = P.vector("wasm-validate", "--enable-exceptions")
            .plusAll(this.target.wabtFlags())
            .plus(path.toString());

        runInheritIo(command);
    }
}