package org.shedlang.compiler.shidtopatio;

import org.shedlang.compiler.patio.PatioFunctionDeclaration;
import org.shedlang.compiler.patio.PatioLocal;
import org.shedlang.compiler.patio.instructions.PatioCall;
import org.shedlang.compiler.patio.instructions.PatioMemoryStore;
import org.shedlang.compiler.patio.operands.*;
import org.shedlang.compiler.patio.types.*;
import org.shedlang.compiler.patio.values.PatioDataPtrValueNull;
import org.shedlang.compiler.patio.values.PatioIntValueByteSizeOf;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.sources.Sources;
import org.shedlang.compiler.util.Lists;

import java.util.List;

public class ShidToPatioGc {
    private ShidToPatioGc() {
    }

    public static final PatioType RUNTIME_TYPE_TAG_TYPE = PatioTypes.INT32;

    public static PatioDataLayout gcObjectDataLayout(PatioDataLayout valueLayout) {
        return PatioDataLayouts.struct(List.of(
            PatioDataLayouts.struct(List.of(
                PatioDataLayouts.DATA_PTR, // next
                PatioDataLayouts.INT32, // mark
                PatioDataLayouts.fromType(RUNTIME_TYPE_TAG_TYPE) // runtime type tag
            )),
            valueLayout
        ));
    }

    public static List<Integer> runtimeTypeTagOffset() {
        return List.of(0, 2);
    }

    public static List<Integer> gcObjectOffset(List<Integer> valueOffset) {
        return Lists.concat(List.of(1), valueOffset);
    }

    /**
     * Creates a data value for GC objects with no GC children that should never
     * be garbage collected. Such a value does not need to be added to the set
     * of GC roots.
     */
    public static PatioDataValue gcObjectStaticDataValue(PatioDataValue value) {
        return PatioDataValues.struct(List.of(
            PatioDataValues.struct(List.of(
                PatioDataValues.int_(0),
                // Static values are always marked: see gc.c
                PatioDataValues.int32(0b110),
                PatioDataValues.int32(0)
            )),
            value
        ));
    }

    // == Roots ==

    public static final PatioType GC_ROOT_STACK_INDEX_TYPE = PatioTypes.INT;

    private static final PatioFunctionDeclaration GET_GC_ROOT_STACK_INDEX = new PatioFunctionDeclaration(
        "shed__runtime__gc_get_root_stack_index",
        List.of(),
        List.of(GC_ROOT_STACK_INDEX_TYPE),
        Sources.RUNTIME
    );

    public static ShidToPatioContext getGcRootStackIndex(
        PatioLocalRef destination,
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(GET_GC_ROOT_STACK_INDEX)
            .addInstruction(new PatioCall(
                List.of(destination),
                GET_GC_ROOT_STACK_INDEX.name(),
                List.of(),
                source
            ));
    }

    private static final PatioFunctionDeclaration SET_GC_ROOT_STACK_INDEX = new PatioFunctionDeclaration(
        "shed__runtime__gc_set_root_stack_index",
        List.of(PatioLocal.param("gc_root_stack_index", GC_ROOT_STACK_INDEX_TYPE)),
        List.of(),
        Sources.RUNTIME
    );

    public static ShidToPatioContext setGcRootStackIndex(
        PatioOperand gcRootIndex,
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(SET_GC_ROOT_STACK_INDEX)
            .addInstruction(new PatioCall(
                List.of(),
                SET_GC_ROOT_STACK_INDEX.name(),
                List.of(gcRootIndex),
                source
            ));
    }

    private static final PatioFunctionDeclaration PUSH_GC_ROOT = new PatioFunctionDeclaration(
        "shed__runtime__gc_push_root",
        List.of(PatioLocal.param("gc_root", ShidToPatioTypes.VALUE_TYPE)),
        List.of(),
        Sources.RUNTIME
    );

    public static ShidToPatioContext pushGcRoot(
        PatioOperand escapee,
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(PUSH_GC_ROOT)
            .addInstruction(new PatioCall(
                List.of(),
                PUSH_GC_ROOT.name(),
                List.of(escapee),
                source
            ));
    }

    // == Initialization ==

    private static final PatioFunctionDeclaration INIT = new PatioFunctionDeclaration(
        "shed__runtime__gc_init",
        List.of(),
        List.of(),
        Sources.RUNTIME
    );

    public static ShidToPatioContext init(
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(INIT)
            .addInstruction(new PatioCall(
                List.of(),
                INIT.name(),
                List.of(),
                source
            ));
    }

    // == Allocation ==

    private static final PatioFunctionDeclaration ALLOC = new PatioFunctionDeclaration(
        "shed__runtime__gc_alloc",
        List.of(
            PatioLocal.param("size", PatioTypes.INT),
            PatioLocal.param("runtime_type_tag", PatioTypes.INT32)
        ),
        List.of(PatioTypes.DATA_PTR),
        Sources.RUNTIME
    );

    public static ShidToPatioContext alloc(
        PatioLocalRef destination,
        RuntimeType runtimeType,
        Source source,
        ShidToPatioContext context
    ) {
        var layout = runtimeType.layout();
        var childrenOffset = runtimeType.gcChildrenOffset();
        var childrenCount = runtimeType.gcChildrenCount();

        var runtimeTypeTagResult = context.runtimeTypeTag(runtimeType);
        var runtimeTypeTag = runtimeTypeTagResult.value();
        context = runtimeTypeTagResult.context();

        context = context.ensureCompilationUnitElement(ALLOC);

        var size = new PatioIntConst(new PatioIntValueByteSizeOf(layout));

        context = context.addInstruction(new PatioCall(
            List.of(destination),
            ALLOC.name(),
            List.of(size, PatioInt32Const.of(runtimeTypeTag)),
            source
        ));

        for (var childIndex = 0; childIndex < childrenCount; childIndex++) {
            var childOffset = Lists.concat(childrenOffset, List.of(childIndex));
            var childDataLayout = PatioOffsets.offsetDataLayout(layout, childOffset);

            var zero = switch (childDataLayout) {
                case PatioDataLayoutDataPtr patioDataLayoutDataPtr ->
                    new PatioDataPtrConst(new PatioDataPtrValueNull());

                case PatioDataLayoutInt dataLayoutInt ->
                    PatioIntConst.of(0);

                default ->
                    // TODO: better error
                    throw new IllegalArgumentException("unsupported child");
            };

            context = context.addInstruction(new PatioMemoryStore(
                destination,
                layout,
                childOffset,
                zero,
                source
            ));
        }

        return context;
    }

    // == Collection  ==

    private static final PatioFunctionDeclaration COLLECT = new PatioFunctionDeclaration(
        "shed__runtime__gc_collect",
        List.of(),
        List.of(),
        Sources.RUNTIME
    );

    public static ShidToPatioContext collect(
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(COLLECT)
            .addInstruction(new PatioCall(
                List.of(),
                COLLECT.name(),
                List.of(),
                source
            ));
    }

    // == Save/Restore ==

    static ShidToPatioContext saveGcRootStack(Source source, ShidToPatioContext context) {
        var gcRootStackIndexResult = context.allocateVal(
            ShidToPatioLocalSpace.GC_ROOT_STACK_INDEX,
            ShidToPatioGc.GC_ROOT_STACK_INDEX_TYPE
        );
        var gcRootStackIndex = gcRootStackIndexResult.value();
        context = gcRootStackIndexResult.context();

        context = ShidToPatioGc.getGcRootStackIndex(
            gcRootStackIndex,
            source,
            context
        );

        context = context.stashGcRootStackIndex(gcRootStackIndex);

        return context;
    }

    static ShidToPatioContext restoreGcRootStack(
        List<PatioOperand> escapees,
        Source source,
        ShidToPatioContext context
    ) {
        context = ShidToPatioGc.setGcRootStackIndex(
            context.stashedGcRootStackIndex(),
            source,
            context
        );

        for (var escapee : escapees) {
            context = ShidToPatioGc.pushGcRoot(escapee, source, context);
        }

        return context;
    }

    // == Clear ==

    private static final PatioFunctionDeclaration CLEAR = new PatioFunctionDeclaration(
        "shed__runtime__gc_clear",
        List.of(),
        List.of(),
        Sources.RUNTIME
    );

    public static ShidToPatioContext clear(
        Source source,
        ShidToPatioContext context
    ) {
        return context
            .ensureCompilationUnitElement(CLEAR)
            .addInstruction(new PatioCall(
                List.of(),
                CLEAR.name(),
                List.of(),
                source
            ));
    }
}