package org.shedlang.compiler.shidtopatio;

import org.shedlang.compiler.ast.typed.FreeVariables;
import org.shedlang.compiler.patio.instructions.*;
import org.shedlang.compiler.patio.operands.PatioIntConst;
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.PatioTypes;
import org.shedlang.compiler.shid.instructions.ShidDefineStruct;
import org.shedlang.compiler.shid.instructions.ShidParam;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.types.SimpleStructType;
import org.shedlang.compiler.types.StructTypeInfo;
import org.shedlang.compiler.util.Lists;

import java.util.List;

public class ShidToPatioStructs {
    private ShidToPatioStructs() {
    }

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

    private static final List<Integer> TAG_OFFSET
        = ShidToPatioGc.runtimeTypeTagOffset();

    private static StructRuntimeType structRuntimeType(
        SimpleStructType type,
        ShidToPatioContext context
    ) {
        var structTypeInfo = context.structTypeInfo(type);
        return structRuntimeType(type, structTypeInfo);
    }

    public static StructRuntimeType structRuntimeType(
        SimpleStructType type,
        StructTypeInfo structTypeInfo
    ) {
        return new StructRuntimeType(type, structTypeInfo.fieldCount());
    }

    public record StructRuntimeType(SimpleStructType type, 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> tagOffset() {
            return TAG_OFFSET;
        }

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

    public static ShidToPatioContext createType(
        PatioLocalRef destination,
        SimpleStructType type,
        List<ShidDefineStruct.Field> fields,
        Source source,
        ShidToPatioContext context
    ) {
        var structTypeInfo = context.structTypeInfo(type);

        var params = fields.stream()
            .map(field -> new ShidParam(field.name(), field.source()))
            .toList();

        return ShidToPatioFunctions.create(
            destination,
            type.name(),
            structTypeInfo.fieldsArePositional() ? params : List.of(),
            structTypeInfo.fieldsArePositional() ? List.of() : params,
            FreeVariables.empty(),
            bodyContext -> {
                var valuePointerResult = bodyContext.allocateVal(
                    ShidToPatioLocalSpace.STRUCT_VALUE_POINTER,
                    PatioTypes.DATA_PTR
                );
                var valuePointer = valuePointerResult.value();
                bodyContext = valuePointerResult.context();

                var structRuntimeType = structRuntimeType(type, context);

                bodyContext = ShidToPatioGc.alloc(
                    valuePointer,
                    structRuntimeType,
                    source,
                    bodyContext
                );

                bodyContext = bodyContext
                    .addInstruction(new PatioMemoryStore(
                        valuePointer,
                        structRuntimeType.layout(),
                        structRuntimeType.tagOffset(),
                        PatioIntConst.of(context.structTag(type)),
                        source
                    ));

                for (var fieldIndex = 0; fieldIndex < fields.size(); fieldIndex++) {
                    var field = fields.get(fieldIndex);

                    bodyContext = bodyContext
                        .addInstruction(new PatioMemoryStore(
                            valuePointer,
                            structRuntimeType.layout(),
                            structRuntimeType.fieldOffset(fieldIndex),
                            bodyContext.shidLocalNameToLocal(field.name()),
                            source
                        ));
                }

                var valueResult = bodyContext.allocateVal(
                    ShidToPatioLocalSpace.STRUCT_VALUE,
                    ShidToPatioTypes.VALUE_TYPE
                );
                var value = valueResult.value();
                bodyContext = valueResult.context();

                bodyContext = bodyContext.addInstruction(new PatioDataPtrToInt(
                    value,
                    valuePointer,
                    source
                ));

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

                return bodyContext;
            },
            source,
            context
        );
    }

    public static ShidToPatioContext fieldLoad(
        PatioLocalRef destination,
        PatioOperand receiver,
        SimpleStructType type,
        String fieldName,
        Source source,
        ShidToPatioContext context
    ) {
        var valuePointerResult = context.allocateVal(
            ShidToPatioLocalSpace.STRUCT_VALUE_POINTER,
            PatioTypes.DATA_PTR
        );
        var valuePointer = valuePointerResult.value();
        context = valuePointerResult.context();

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

        var structRuntimeType = structRuntimeType(type, context);

        var structTypeInfo = context.structTypeInfo(type);
        var fieldIndex = structTypeInfo.fieldIndex(fieldName);

        context = context.addInstruction(new PatioMemoryLoad(
            destination,
            valuePointer,
            structRuntimeType.layout(),
            structRuntimeType.fieldOffset(fieldIndex),
            source
        ));

        return context;
    }

    public static ShidToPatioContext tagLoad(
        PatioLocalRef destination,
        PatioOperand receiver,
        Source source,
        ShidToPatioContext context
    ) {
        var valuePointerResult = context.allocateVal(
            ShidToPatioLocalSpace.STRUCT_VALUE_POINTER,
            PatioTypes.DATA_PTR
        );
        var valuePointer = valuePointerResult.value();
        context = valuePointerResult.context();

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

        var structDataLayout = unknownStructDataLayout();

        context = context.addInstruction(new PatioMemoryLoad(
            destination,
            valuePointer,
            structDataLayout,
            TAG_OFFSET,
            source
        ));

        return context;
    }
}