package org.shedlang.compiler.shidtopatio;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.shedlang.compiler.ast.typed.FreeVariables;
import org.shedlang.compiler.patio.types.PatioTypes;
import org.shedlang.compiler.patio.instructions.Patio;
import org.shedlang.compiler.sources.Sources;
import org.shedlang.compiler.shid.instructions.ShidDefineFunction;
import org.shedlang.compiler.shid.instructions.Shid;
import org.shedlang.compiler.testing.snapshots.SnapshotResolver;
import org.shedlang.compiler.testing.snapshots.Snapshotter;

import java.io.IOException;
import java.util.List;

import static org.shedlang.compiler.shidtopatio.ShidToPatioTesting.serialize;

@ExtendWith({SnapshotResolver.class})
public class ShidToPatioDefineFunctionTests {
    @Test
    public void canDefineMinimalFunction(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(),
            List.of(),
            FreeVariables.empty(),
            Shid.block(List.of(
                Shid.pushUnit(),
                Shid.return_()
            )),
            ShidDefineFunction.Variant.FUNCTION,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY;

        var result = ShidToPatio.generateInstruction(instruction, context);

        snapshotter.assertSnapshot(serialize(result));
    }

    @Test
    public void immutableFreeVariablesAreStoredInAndLoadedFromEnvironment(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(),
            List.of(),
            FreeVariables.builder().plus("a", Sources.NULL).plus("b", Sources.NULL).build(),
            Shid.block(List.of(
                Shid.localLoad("a"),
                Shid.return_()
            )),
            ShidDefineFunction.Variant.FUNCTION,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY
            .defineShidParam("a")
            .defineShidParam("b");

        var result = ShidToPatio.generateInstruction(instruction, context);

        snapshotter.assertSnapshot(serialize(result));
    }

    @Test
    public void canCaptureFreeVariablesThatAreSetAfterDefinition(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(),
            List.of(),
            FreeVariables.builder().plus("a", Sources.NULL).plus("b", Sources.NULL).build(),
            Shid.block(List.of(
                Shid.localLoad("a"),
                Shid.return_()
            )),
            ShidDefineFunction.Variant.FUNCTION,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY
            .declareShidVal("a", ShidToPatioTypes.VALUE_TYPE).context()
            .declareShidVal("b", ShidToPatioTypes.VALUE_TYPE).context();

        context = ShidToPatio.generateInstruction(instruction, context);

        // We make sure there are some other instructions before setting the
        // free variables so we can distinguish the free variable being copied
        // into the closure upon closure creation (i.e. too soon) versus after
        // the free variable is set.
        context = context.addInstruction(Patio.localStore(Patio.localRef("a"), Patio.intConst(42)));
        context = context.initShidLocal("a");
        context = context.addInstruction(Patio.localStore(Patio.localRef("b"), Patio.intConst(47)));
        context = context.initShidLocal("b");

        snapshotter.assertSnapshot(serialize(context));
    }

    @Test
    public void pointersToFreeMutableVariablesAreStoredInAndLoadedFromEnvironment(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(),
            List.of(),
            FreeVariables.builder()
                .plus("a", Sources.NULL)
                .plus("b", Sources.NULL)
                .build(),
            Shid.block(List.of(
                Shid.localLoad("a"),
                Shid.return_()
            )),
            ShidDefineFunction.Variant.FUNCTION,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY
            .declareShidVal("a", PatioTypes.DATA_PTR).context()
            .initShidLocal("a")
            .declareShidVal("b", PatioTypes.DATA_PTR).context()
            .initShidLocal("b");

        var result = ShidToPatio.generateInstruction(instruction, context);

        snapshotter.assertSnapshot(serialize(result));
    }

    @Test
    public void positionalAndNamedParamsAreDefined(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(Shid.param("param1"), Shid.param("param2")),
            // Define named params in non-lexical order
            List.of(Shid.param("paramB"), Shid.param("paramC"), Shid.param("paramA")),
            FreeVariables.builder().plus("a", Sources.NULL).plus("b", Sources.NULL).build(),
            Shid.block(List.of(
                Shid.localLoad("a"),
                Shid.localLoad("param1"),
                Shid.intAdd(),
                Shid.localLoad("paramA"),
                Shid.intAdd(),
                Shid.return_()
            )),
            ShidDefineFunction.Variant.FUNCTION,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY
            .defineShidParam("a")
            .defineShidParam("b");

        var result = ShidToPatio.generateInstruction(instruction, context);

        snapshotter.assertSnapshot(serialize(result));
    }

    @Test
    public void effectHandlerMethodSwapsEffectHandlerInBody(Snapshotter snapshotter) throws IOException {
        var instruction = new ShidDefineFunction(
            "f",
            List.of(),
            List.of(),
            FreeVariables.empty(),
            Shid.block(List.of(
                Shid.pushUnit(),
                Shid.resume()
            )),
            ShidDefineFunction.Variant.EFFECT_HANDLER_METHOD,
            Sources.NULL
        );
        var context = ShidToPatioContext.EMPTY;

        var result = ShidToPatio.generateInstruction(instruction, context);

        snapshotter.assertSnapshot(serialize(result));
    }
}