package org.shedlang.compiler.shidtopatio;

import org.shedlang.compiler.patio.PatioFunctionDeclaration;
import org.shedlang.compiler.patio.PatioVisibility;
import org.shedlang.compiler.patio.instructions.*;
import org.shedlang.compiler.patio.operands.PatioBoolConst;
import org.shedlang.compiler.patio.operands.PatioDataPtrConst;
import org.shedlang.compiler.patio.operands.PatioLocalRef;
import org.shedlang.compiler.patio.operands.PatioOperand;
import org.shedlang.compiler.patio.types.*;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.types.NamespaceName;
import org.shedlang.compiler.util.Lists;

import java.util.List;
import java.util.function.Function;

public class ShidToPatioNamespaces {
    public static final PatioType NAMESPACE_VALUE_TYPE = ShidToPatioTypes.VALUE_TYPE;

    private ShidToPatioNamespaces() {}

    private static NamespaceRuntimeType namespaceRuntimeType(
        NamespaceName namespaceName,
        ShidToPatioContext context
    ) {
        var namespaceTypeInfo = context.namespaceTypeInfo(namespaceName);
        return new NamespaceRuntimeType(namespaceTypeInfo.fieldCount());
    }

    private record NamespaceRuntimeType(int fieldCount) implements RuntimeType {
        @Override
        public PatioDataLayout layout() {
            return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(
                Lists.repeat(
                    PatioDataLayouts.INT,
                    fieldCount
                )
            ));
        }

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

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

        @Override
        public int gcChildrenCount() {
            return fieldCount;
        }

        public List<Integer> fieldOffset(int fieldIndex) {
            return ShidToPatioGc.gcObjectOffset(List.of(fieldIndex));
        }
    }

    public static ShidToPatioContext define(
        NamespaceName namespaceName,
        Function<ShidToPatioContext, ShidToPatioContext> generateInitBody,
        Source source,
        ShidToPatioContext context
    ) {
        var namespaceRuntimeType = namespaceRuntimeType(namespaceName, context);

        context = context.addDataElementMutable(
            globalNameNamespaceValue(namespaceName),
            PatioVisibility.PRIVATE,
            PatioDataValues.dataPtrNull(),
            source
        );

        context = context.addDataElementMutable(
            globalNameNamespaceIsInited(namespaceName),
            PatioVisibility.PRIVATE,
            PatioDataValues.bool(false),
            source
        );

        context = context.createFunction(
            globalNameNamespaceInit(namespaceName),
            PatioVisibility.PRIVATE,
            List.of(),
            source,
            initContext -> {
                var namespaceValueResult = initContext.allocateVal(
                    ShidToPatioLocalSpace.NAMESPACE_VALUE,
                    PatioTypes.DATA_PTR
                );
                var namespaceValue = namespaceValueResult.value();
                initContext = namespaceValueResult.context();

                initContext = ShidToPatioGc.alloc(
                    namespaceValue,
                    namespaceRuntimeType,
                    source,
                    initContext
                );

                initContext = initContext.addInstruction(new PatioMemoryStore(
                    globalNameNamespaceValueOperand(namespaceName),
                    PatioDataLayouts.DATA_PTR,
                    List.of(),
                    namespaceValue,
                    source
                ));

                initContext = generateInitBody.apply(initContext);

                initContext = initContext.addInstruction(new PatioMemoryStore(
                    globalNameNamespaceIsInitedOperand(namespaceName),
                    PatioDataLayouts.BOOL,
                    List.of(),
                    new PatioBoolConst(true),
                    source
                ));

                initContext = initContext.addInstruction(new PatioReturn(List.of(), source));

                return initContext;
            }
        );

        context = context.createFunction(
            globalNameNamespaceImport(namespaceName),
            PatioVisibility.PUBLIC,
            List.of(NAMESPACE_VALUE_TYPE),
            source,
            importContext -> {
                var isInitedResult = importContext.allocateVal(
                    ShidToPatioLocalSpace.TEMPORARY,
                    PatioTypes.BOOL
                );
                var isInited = isInitedResult.value();
                importContext = isInitedResult.context();

                importContext = importContext.addInstruction(new PatioMemoryLoad(
                    isInited,
                    globalNameNamespaceIsInitedOperand(namespaceName),
                    PatioDataLayouts.BOOL,
                    List.of(),
                    source
                ));

                importContext = importContext.addInstruction(new PatioIf(
                    List.of(),
                    isInited,
                    new PatioBlock(
                        List.of(),
                        List.of(),
                        source
                    ),
                    new PatioBlock(
                        List.of(
                            new PatioCall(
                                List.of(),
                                globalNameNamespaceInit(namespaceName),
                                List.of(),
                                source
                            )
                        ),
                        List.of(),
                        source
                    ),
                    source
                ));

                var namespaceValuePointerResult = importContext.allocateVal(
                    ShidToPatioLocalSpace.NAMESPACE_POINTER,
                    NAMESPACE_VALUE_TYPE
                );
                var namespaceValuePointer = namespaceValuePointerResult.value();
                importContext = namespaceValuePointerResult.context();

                importContext = importContext.addInstruction(new PatioMemoryLoad(
                    namespaceValuePointer,
                    globalNameNamespaceValueOperand(namespaceName),
                    PatioDataLayouts.DATA_PTR,
                    List.of(),
                    source
                ));

                var namespaceValueResult = importContext.allocateVal(
                    ShidToPatioLocalSpace.NAMESPACE_VALUE,
                    NAMESPACE_VALUE_TYPE
                );
                var namespaceValue = namespaceValueResult.value();
                importContext = namespaceValueResult.context();

                importContext = importContext.addInstruction(new PatioDataPtrToInt(
                    namespaceValue,
                    namespaceValuePointer,
                    source
                ));

                importContext = importContext.addInstruction(new PatioReturn(
                    List.of(namespaceValue),
                    source
                ));

                return importContext;
            }
        );

        return context;
    }

    public static ShidToPatioContext importValue(
        PatioLocalRef destination,
        NamespaceName namespaceName,
        Source source,
        ShidToPatioContext context
    ) {
        context = context.addCompilationUnitElement(new PatioFunctionDeclaration(
            globalNameNamespaceImport(namespaceName),
            List.of(),
            List.of(NAMESPACE_VALUE_TYPE),
            source
        ));

        context = context.addInstruction(new PatioCall(
            List.of(destination),
            globalNameNamespaceImport(namespaceName),
            List.of(),
            source
        ));

        return context;
    }

    public static ShidToPatioContext fieldLoad(
        PatioLocalRef destination,
        PatioOperand receiver,
        NamespaceName namespaceName,
        String fieldName,
        Source source,
        ShidToPatioContext context
    ) {
        var namespaceTypeInfo = context.namespaceTypeInfo(namespaceName);
        var namespaceRuntimeType = namespaceRuntimeType(namespaceName, context);

        var namespacePointerResult = context.allocateVal(
            ShidToPatioLocalSpace.NAMESPACE_POINTER,
            namespaceNameToPatioIdentifier(namespaceName),
            PatioTypes.DATA_PTR
        );
        var namespacePointer = namespacePointerResult.value();
        context = namespacePointerResult.context();

        context = context.addInstruction(new PatioIntToDataPtr(
            namespacePointer,
            receiver,
            source
        ));

        context = context.addInstruction(new PatioMemoryLoad(
            destination,
            namespacePointer,
            namespaceRuntimeType.layout(),
            namespaceRuntimeType.fieldOffset(namespaceTypeInfo.fieldIndex(fieldName)),
            source
        ));

        return context;
    }

    public static ShidToPatioContext fieldStore(
        NamespaceName namespaceName,
        String fieldName,
        PatioOperand fieldValue,
        Source source,
        ShidToPatioContext context
    ) {
        var namespaceTypeInfo = context.namespaceTypeInfo(namespaceName);
        var namespaceRuntimeType = namespaceRuntimeType(namespaceName, context);

        var namespaceValueResult = context.allocateVal(
            ShidToPatioLocalSpace.NAMESPACE_VALUE,
            PatioTypes.DATA_PTR
        );
        var namespaceValue = namespaceValueResult.value();
        context = namespaceValueResult.context();

        context = context.addInstruction(new PatioMemoryLoad(
            namespaceValue,
            globalNameNamespaceValueOperand(namespaceName),
            PatioDataLayouts.DATA_PTR,
            List.of(),
            source
        ));

        context = context.addInstruction(new PatioMemoryStore(
            namespaceValue,
            namespaceRuntimeType.layout(),
            namespaceRuntimeType.fieldOffset(namespaceTypeInfo.fieldIndex(fieldName)),
            fieldValue,
            source
        ));

        return context;
    }

    private static PatioOperand globalNameNamespaceValueOperand(NamespaceName namespaceName) {
        return PatioDataPtrConst.dataElement(globalNameNamespaceValue(namespaceName));
    }

    // TODO: allocate names through context
    private static String globalNameNamespaceValue(NamespaceName namespaceName) {
        return "shed__namespace__" + namespaceNameToPatioIdentifier(namespaceName);
    }

    private static String globalNameNamespaceInit(NamespaceName namespaceName) {
        return "shed__namespaceInit__" + namespaceNameToPatioIdentifier(namespaceName);
    }

    private static PatioOperand globalNameNamespaceIsInitedOperand(NamespaceName namespaceName) {
        return PatioDataPtrConst.dataElement(globalNameNamespaceIsInited(namespaceName));
    }

    private static String globalNameNamespaceIsInited(NamespaceName namespaceName) {
        return "shed__namespaceIsInited__" + namespaceNameToPatioIdentifier(namespaceName);
    }

    private static String globalNameNamespaceImport(NamespaceName namespaceName) {
        return "shed__namespaceImport__" + namespaceNameToPatioIdentifier(namespaceName);
    }

    public static String namespaceNameToPatioIdentifier(NamespaceName namespaceName) {
        return String.join("_", namespaceName.parts());
    }
}