package org.shedlang.compiler.shidtopatio;

import org.shedlang.compiler.ast.typed.FreeVariable;
import org.shedlang.compiler.ast.typed.FreeVariables;
import org.shedlang.compiler.patio.PatioFunctionDeclaration;
import org.shedlang.compiler.patio.PatioLocal;
import org.shedlang.compiler.patio.PatioVisibility;
import org.shedlang.compiler.patio.instructions.*;
import org.shedlang.compiler.patio.operands.PatioFunctionPtrConst;
import org.shedlang.compiler.patio.operands.PatioLocalRef;
import org.shedlang.compiler.patio.operands.PatioOperand;
import org.shedlang.compiler.patio.types.PatioDataLayout;
import org.shedlang.compiler.patio.types.PatioDataLayouts;
import org.shedlang.compiler.patio.types.PatioType;
import org.shedlang.compiler.patio.types.PatioTypes;
import org.shedlang.compiler.shid.instructions.ShidParam;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.types.NamespaceName;
import org.shedlang.compiler.util.P;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Stream;

import static org.shedlang.compiler.util.Iterables.sortedBy;

class ShidToPatioClosures {
    public static final PatioType CLOSURE_POINTER_TYPE = PatioTypes.DATA_PTR;
    public static final PatioType CLOSURE_TYPE = ShidToPatioTypes.VALUE_TYPE;

    private ShidToPatioClosures() {}

    public static ClosureRuntimeType closureRuntimeType(
        PatioDataLayout additional,
        FreeVariables freeVariables,
        ShidToPatioContext context
    ) {
        return new ClosureRuntimeType(
            additional,
            freeVariables.stream()
                .map(freeVariable -> PatioDataLayouts.fromType(
                    context.shidLocalType(freeVariable.name())
                ))
                .toList()
        );
    }

    public record ClosureRuntimeType(
        PatioDataLayout additional,
        List<PatioDataLayout> freeVariables
    ) implements RuntimeType {
        @Override
        public PatioDataLayout layout() {
            return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(List.of(
                PatioDataLayouts.FUNCTION_PTR,
                additional,
                PatioDataLayouts.struct(freeVariables)
            )));
        }

        private static final List<Integer> FUNCTION_PTR_OFFSET
            = ShidToPatioGc.gcObjectOffset(List.of(0));

        public List<Integer> functionPtrOffset() {
            return FUNCTION_PTR_OFFSET;
        }

        private static final List<Integer> GC_CHILDREN_OFFSET =
            ShidToPatioGc.gcObjectOffset(List.of(2));

        @Override
        public List<Integer> gcChildrenOffset() {
            return GC_CHILDREN_OFFSET;
        }

        @Override
        public int gcChildrenCount() {
            return freeVariables.size();
        }

        public List<Integer> freeVariableOffset(
            FreeVariable freeVariable
        ) {
            return ShidToPatioGc.gcObjectOffset(List.of(2, freeVariable.index()));
        }
    }

    private static PatioDataLayout closureDataLayoutForCall() {
        return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(List.of(
            PatioDataLayouts.FUNCTION_PTR
        )));
    }

    static List<Integer> additionalOffset(List<Integer> offset) {
        return ShidToPatioGc.gcObjectOffset(P.vector(1).plusAll(offset));
    }

    static ShidToPatioContext defineParams(
        List<ShidParam> positionalParams,
        List<ShidParam> namedParams,
        ShidToPatioContext bodyContext
    ) {
        var shidParams = generateShidParamList(positionalParams, namedParams);

        for (var shidParam : shidParams) {
            bodyContext = bodyContext.defineShidParam(shidParam.name());
        }

        return bodyContext;
    }

    private static List<ShidParam> generateShidParamList(
        List<ShidParam> positionalParams,
        List<ShidParam> namedParams
    ) {
        var orderedNamedParams = sortedBy(namedParams, param -> param.name());
        return Stream.concat(positionalParams.stream(), orderedNamedParams.stream()).toList();
    }

    static ShidToPatioContext create(
        PatioLocalRef destination,
        ShidToPatioCompilationUnitElementSpace functionElementSpace,
        String name,
        List<ShidParam> positionalParams,
        List<ShidParam> namedParams,
        ClosureRuntimeType closureRuntimeType,
        BiFunction<PatioLocalRef, ShidToPatioContext, ShidToPatioContext> storeAdditional,
        FreeVariables freeVariables,
        BiFunction<PatioLocalRef, ShidToPatioContext, ShidToPatioContext> generateBody,
        Source source,
        ShidToPatioContext context
    ) {
        var fullNameResult = context.allocateCompilationUnitElementName(
            functionElementSpace,
            name
        );
        var fullName = fullNameResult.value();
        context = fullNameResult.context();

        var bodyContextParent = context;
        context = context.createFunction(fullName, PatioVisibility.PRIVATE, List.of(ShidToPatioTypes.VALUE_TYPE), source, bodyContext -> {
            var closurePointerResult = bodyContext.allocateParam(
                ShidToPatioLocalSpace.CLOSURE_POINTER,
                ShidToPatioClosures.CLOSURE_POINTER_TYPE
            );
            var closurePointer = closurePointerResult.value();
            bodyContext = closurePointerResult.context();

            bodyContext = defineParams(positionalParams, namedParams, bodyContext);

            bodyContext = loadFreeVariables(
                freeVariables,
                closurePointer,
                closureRuntimeType,
                source,
                bodyContext,
                bodyContextParent
            );

            bodyContext = ShidToPatioGc.saveGcRootStack(source, bodyContext);

            bodyContext = generateBody.apply(closurePointer, bodyContext);

            return bodyContext;
        });

        context = ShidToPatioClosures.alloc(
            destination,
            name,
            fullName,
            closureRuntimeType,
            storeAdditional,
            freeVariables,
            source,
            context
        );

        return context;
    }

    public static ShidToPatioContext createNative(
        PatioLocalRef destination,
        NamespaceName namespaceName,
        String name,
        List<ShidParam> positionalParams,
        List<ShidParam> namedParams,
        ClosureRuntimeType closureRuntimeType,
        Source source,
        ShidToPatioContext context
    ) {
        var fullName = String.format(
            "shed__runtime__namespace__%s__%s",
            String.join("__", namespaceName.parts()),
            name
        );

        context = context.addCompilationUnitElement(new PatioFunctionDeclaration(
            fullName,
            generateShidParamList(positionalParams, namedParams).stream()
                .map(param -> PatioLocal.param(param.name(), ShidToPatioTypes.VALUE_TYPE))
                .toList(),
            List.of(ShidToPatioTypes.VALUE_TYPE),
            source
        ));

        context = alloc(
            destination,
            name,
            fullName,
            closureRuntimeType,
            (closurePointer, storeContext) -> storeContext,
            FreeVariables.empty(),
            source,
            context
        );

        return context;
    }


    static ShidToPatioContext alloc(
        PatioLocalRef destination,
        String name,
        String functionName,
        ClosureRuntimeType closureRuntimeType,
        BiFunction<PatioLocalRef, ShidToPatioContext, ShidToPatioContext> storeAdditional,
        FreeVariables freeVariables,
        Source source,
        ShidToPatioContext context
    ) {
        var closurePointerResult = context.allocateVal(
            ShidToPatioLocalSpace.CLOSURE_POINTER,
            name,
            PatioTypes.DATA_PTR
        );
        var closurePointer = closurePointerResult.value();
        context = closurePointerResult.context();

        context = ShidToPatioGc.alloc(
            closurePointer,
            closureRuntimeType,
            source,
            context
        );

        context = context.addInstruction(new PatioMemoryStore(
            closurePointer,
            closureRuntimeType.layout(),
            closureRuntimeType.functionPtrOffset(),
            new PatioFunctionPtrConst(functionName),
            source
        ));

        context = storeAdditional.apply(closurePointer, context);

        context = storeFreeVariables(
            freeVariables,
            closurePointer,
            closureRuntimeType,
            source,
            context
        );

        context = context.addInstruction(new PatioDataPtrToInt(
            destination,
            closurePointer,
            source
        ));

        return context;
    }

    static ShidToPatioContext storeFreeVariables(
        FreeVariables freeVariables,
        PatioLocalRef closurePointer,
        ClosureRuntimeType closureRuntimeType,
        Source source,
        ShidToPatioContext context
    ) {
        for (var freeVariable : freeVariables) {
            context = context.onceShidLocalDefined(
                freeVariable.name(),
                c -> c.addInstruction(new PatioMemoryStore(
                    closurePointer,
                    closureRuntimeType.layout(),
                    closureRuntimeType.freeVariableOffset(freeVariable),
                    c.shidLocalNameToLocal(freeVariable.name()),
                    source
                ))
            );
        }
        return context;
    }

    static ShidToPatioContext loadFreeVariables(
        FreeVariables freeVariables,
        PatioLocalRef closurePointer,
        ClosureRuntimeType closureRuntimeType,
        Source source,
        ShidToPatioContext bodyContext,
        ShidToPatioContext bodyContextParent
    ) {
        for (var freeVariable : freeVariables) {
            var freeVariableLocalResult = bodyContext.declareShidVal(
                freeVariable.name(),
                bodyContextParent.shidLocalType(freeVariable.name())
            );
            var freeVariableLocal = freeVariableLocalResult.value();
            bodyContext = freeVariableLocalResult.context();

            bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                freeVariableLocal,
                closurePointer,
                closureRuntimeType.layout(),
                closureRuntimeType.freeVariableOffset(freeVariable),
                source
            ));

            bodyContext = bodyContext.initShidLocal(freeVariable.name());
        }
        return bodyContext;
    }

    static ShidToPatioContext beforeExit(
        List<PatioOperand> escapees,
        Source source,
        ShidToPatioContext context
    ) {
        return ShidToPatioGc.restoreGcRootStack(escapees, source, context);
    }

    static ShidToPatioContext call(
        PatioLocalRef destination,
        PatioLocalRef closure,
        List<PatioLocalRef> positionalArgs,
        List<Map.Entry<String, PatioLocalRef>> namedArgsByName,
        Source source,
        ShidToPatioContext context
    ) {
        var closurePointerResult = context.allocateVal(
            ShidToPatioLocalSpace.CLOSURE_POINTER,
            CLOSURE_POINTER_TYPE
        );
        var closurePointer = closurePointerResult.value();
        context = closurePointerResult.context();

        context = context.addInstruction(new PatioIntToDataPtr(
            closurePointer,
            closure,
            source
        ));

        var functionPointerResult = context.allocateVal(
            ShidToPatioLocalSpace.CALL_FUNCTION_POINTER,
            PatioTypes.FUNCTION_PTR
        );
        var functionPointer = functionPointerResult.value();
        context = functionPointerResult.context();

        context = context.addInstruction(new PatioMemoryLoad(
            functionPointer,
            closurePointer,
            closureDataLayoutForCall(),
            ClosureRuntimeType.FUNCTION_PTR_OFFSET,
            source
        ));

        var args = new ArrayList<PatioOperand>();

        args.add(closurePointer);

        args.addAll(positionalArgs);

        for (var namedArg : sortedBy(namedArgsByName, x -> x.getKey())) {
            args.add(namedArg.getValue());
        }

        return context
            .addInstruction(new PatioCallPointer(
                List.of(destination),
                functionPointer,
                args,
                source
            ));
    }
}