package org.shedlang.compiler.shid;

import org.pcollections.*;
import org.shedlang.compiler.ShedBuiltins;
import org.shedlang.compiler.ast.typed.*;
import org.shedlang.compiler.shid.instructions.*;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.sources.Sources;
import org.shedlang.compiler.types.*;
import org.shedlang.compiler.util.Iterables;
import org.shedlang.compiler.util.P;

import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.stream.IntStream;
import java.util.stream.Stream;

public class ShidCodeGen {
    private ShidCodeGen() {
    }

    private static Context compileAssign(TypedAssignNode node, Context context) {
        context = compileExpression(node.expression(), context);
        context = compileTarget(node.target(), context);
        return context;
    }

    private static Context compileBinaryOperation(
        TypedBinaryOperationNode node,
        Context context
    ) {
        context = compileExpression(node.left(), context);

        context = compileExpression(node.right(), context);

        context = context.addInstruction(switch (node.operator()) {
            case INT_ADD ->
                new ShidIntAdd(node.source());
            case INT_EQUALS ->
                new ShidIntEquals(node.source());
            case INT_GREATER_THAN ->
                new ShidIntGreaterThan(node.source());
            case INT_GREATER_THAN_OR_EQUAL ->
                new ShidIntGreaterThanOrEqual(node.source());
            case INT_LESS_THAN ->
                new ShidIntLessThan(node.source());
            case INT_LESS_THAN_OR_EQUAL ->
                new ShidIntLessThanOrEqual(node.source());
            case INT_MULTIPLY ->
                new ShidIntMultiply(node.source());
            case INT_NOT_EQUAL ->
                new ShidIntNotEqual(node.source());
            case INT_SUB ->
                new ShidIntSub(node.source());
            case STRING_ADD ->
                new ShidStringAdd(node.source());
        });

        return context;
    }

    static Context compileBlockWithYield(TypedBlockNode node, Context context) {
        for (var statement : node.statements()) {
            context = compileFunctionStatement(statement, context);
        }

        switch (node.behaviour()) {
            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
            }
            case TypedFunctionStatementBehaviour.Continue continue_ -> {
                context = context.addInstruction(new ShidPushUnit(node.source()));
            }
            case TypedFunctionStatementBehaviour.Terminate terminate -> {
            }
        }

        return context;
    }

    static Context compileBlockWithoutYield(TypedBlockNode node, Context context) {
        for (var statement : node.statements()) {
            context = compileFunctionStatement(statement, context);
        }

        switch (node.behaviour()) {
            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
                context = context.addInstruction(new ShidDiscard(node.source()));
            }
            case TypedFunctionStatementBehaviour.Continue continue_ -> {
            }
            case TypedFunctionStatementBehaviour.Terminate terminate -> {
            }
        }

        return context;
    }

    private static Context compileBlockReturn(TypedBlockReturnNode blockReturn, Context context) {
        return compileExpression(blockReturn.expression(), context);
    }

    private static Context compileBoolLiteral(TypedBoolLiteralNode node, Context context) {
        return context.addInstruction(new ShidPushBool(node.value(), node.source()));
    }

    private static Context compileBreak(TypedBreakNode node, Context context) {
        context = compileExpression(node.expression(), context);
        context = context.addInstruction(new ShidEffectHandlerBreak(node.handledEffect(), node.source()));
        return context;
    }

    private static Context compileCall(TypedCallNode node, Context context) {
        var extraArgumentCount = 0;
        if (
            node.receiver() instanceof TypedMemberAccessNode receiverNode &&
                receiverNode.memberKind() == MemberKind.METHOD
        ) {
            extraArgumentCount = 1;
            // TODO: handle not concrete type
            var methodCallReceiverType = (SimpleConcreteType) Types.underlyingType(receiverNode.receiver().type());

            context = context.addInstruction(new ShidMethodLoad(
                methodCallReceiverType,
                receiverNode.memberName(),
                receiverNode.source()
            ));

            context = compileExpression(receiverNode.receiver(), context);
        } else {
            context = compileExpression(node.receiver(), context);
        }

        for (var arg : node.args().positional()) {
            context = compileExpression(arg, context);
        }

        for (var arg : node.args().named()) {
            context = compileExpression(arg.expression(), context);
        }

        var callInstruction = new ShidCall(
            node.args().positional().size() + extraArgumentCount,
            node.args().named().stream()
                .map(arg -> arg.name())
                .toList(),
            node.source()
        );

        context = context.addInstruction(callInstruction);

        return context;
    }

    private static Context compileEffectDefinition(
        TypedEffectDefinitionNode node,
        Context context
    ) {
        var localName = context.shedVariableToLocal(node.effect().name());

        return context
            .addInstruction(new ShidDefineEffect(node.effect(), node.source()))
            .addInstruction(new ShidLocalStore(localName, node.source()));
    }

    public static Context compileExpression(TypedExpressionNode node, Context context) {
        return switch (node) {
            case TypedBinaryOperationNode typedBinaryOperationNode ->
                compileBinaryOperation(typedBinaryOperationNode, context);

            case TypedBoolLiteralNode typedBoolLiteralNode ->
                compileBoolLiteral(typedBoolLiteralNode, context);

            case TypedCallNode typedCallNode ->
                compileCall(typedCallNode, context);

            case TypedFunctionExpressionNode typedFunctionExpressionNode ->
                compileFunctionExpression(typedFunctionExpressionNode, context);

            case TypedIfNode typedIfNode ->
                compileIf(typedIfNode, context);

            case TypedIntLiteralNode typedIntLiteralNode ->
                compileIntLiteral(typedIntLiteralNode, context);

            case TypedMemberAccessNode typedMemberAccessNode ->
                compileMemberAccess(typedMemberAccessNode, context);

            case TypedReferenceNode typedReferenceNode ->
                compileReference(typedReferenceNode, context);

            case TypedStringLiteralNode typedStringLiteralNode ->
                compileStringLiteral(typedStringLiteralNode, context);

            case TypedTryNode typedTryNode ->
                compileTry(typedTryNode, context);

            case TypedTupleNode typedTupleNode ->
                compileTuple(typedTupleNode, context);

            case TypedUnitLiteralNode typedUnitLiteralNode ->
                compileUnitLiteral(typedUnitLiteralNode, context);

            case TypedWhenNode typedWhenNode ->
                compileWhen(typedWhenNode, context);
        };
    }

    private static Context compileExpressionStatement(
        TypedExpressionStatementNode node,
        Context context
    ) {
        context = compileExpression(node.expression(), context);
        context = context.addInstruction(new ShidDiscard(node.source()));
        return context;
    }

    private static Context compileFunction(
        String name,
        TypedParamsNode params,
        TypedFunctionBodyNode bodyNode,
        ShidDefineFunction.Variant variant,
        Source source,
        Context context
    ) {
        context = context.enterBlock(bodyNode.boundVariables(), bodyNode.source());

        var c = context;

        var positionalParams = params.positional().stream()
            .map(param -> new ShidParam(
                c.shedVariableToLocal(param.name()),
                param.source()
            ))
            .toList();

        var namedParams = params.named().stream()
            .map(param -> new ShidParam(
                c.shedVariableToLocal(param.name()),
                param.source()
            ))
            .toList();

        context = compileBlockWithYield(bodyNode.block(), context);

        switch (bodyNode.behaviour()) {
            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
                context = context.addInstruction(new ShidReturn(bodyNode.source()));
            }
            case TypedFunctionStatementBehaviour.Continue continue_ -> {
                context = context.addInstruction(new ShidReturn(bodyNode.source()));
            }
            case TypedFunctionStatementBehaviour.Terminate terminate -> {
            }
        }

        var bodyResult = context.exitBlock();
        var body = bodyResult.value;
        context = bodyResult.context;

        var freeVariables = generateFreeVariables(bodyNode.freeVariables(), context);

        var defineFunctionInstruction = new ShidDefineFunction(
            name,
            positionalParams,
            namedParams,
            freeVariables,
            body,
            variant,
            source
        );

        return context.addInstruction(defineFunctionInstruction);
    }

    private static FreeVariables generateFreeVariables(FreeVariables shedFreeVariables, Context context) {
        var freeVariables = FreeVariables.builder();
        for (var shedFreeVariable : shedFreeVariables) {
            var freeVariableLocalName = context.shedVariableToLocal(shedFreeVariable.name());
            for (var freeVariableSource : shedFreeVariable.referenceSources()) {
                freeVariables = freeVariables.plus(
                    freeVariableLocalName,
                    freeVariableSource
                );
            }
        }
        return freeVariables.build();
    }

    private static Context compileFunctionDefinition(
        TypedFunctionDefinitionNode node,
        Context context
    ) {
        context = compileFunction(
            node.name(),
            node.signature().params(),
            node.body(),
            ShidDefineFunction.Variant.FUNCTION,
            node.source(),
            context
        );

        var localName = context.shedVariableToLocal(node.name());
        var storeInstruction = new ShidLocalStore(localName, node.source());

        return context.addInstruction(storeInstruction);
    }

    private static Context compileFunctionExpression(
        TypedFunctionExpressionNode node,
        Context context
    ) {
        return compileFunction(
            // TODO: include source location?
            "anonymous",
            node.signature().params(),
            node.body(),
            ShidDefineFunction.Variant.FUNCTION,
            node.source(),
            context
        );
    }

    private static Context compileFunctionReturn(
        TypedFunctionReturnNode node,
        Context context
    ) {
        context = compileExpression(node.expression(), context);
        context = context.addInstruction(new ShidReturn(node.source()));
        return context;
    }

    public static Context compileFunctionStatement(
        TypedFunctionStatementNode node,
        Context context
    ) {
        return switch (node) {
            case TypedAssignNode assign ->
                compileAssign(assign, context);

            case TypedBlockReturnNode blockReturn ->
                compileBlockReturn(blockReturn, context);

            case TypedBreakNode break_ ->
                compileBreak(break_, context);

            case TypedExpressionStatementNode expressionStatement ->
                compileExpressionStatement(expressionStatement, context);

            case TypedFunctionReturnNode functionReturn ->
                compileFunctionReturn(functionReturn, context);

            case TypedResumeNode resume ->
                compileResume(resume, context);

            case TypedWhileNode while_ ->
                compileWhile(while_, context);
        };
    }

    private static Context compileIf(TypedIfNode node, Context context) {
        return compileIf(
            node.conditionalBranches(),
            node.elseBody(),
            context
        );
    }

    private static Context compileIf(
        List<TypedConditionalBranchNode> conditionalBranches,
        TypedBlockNode elseBody,
        Context context
    ) {
        if (conditionalBranches.isEmpty()) {
            context = context.enterBlock(elseBody.boundVariables(), elseBody.source());
            context = compileBlockWithYield(elseBody, context);

            var bodyResult = context.exitBlock();
            var body = bodyResult.value;
            context = bodyResult.context;

            context = context.addInstruction(body);
        } else {
            var conditionalBranch = conditionalBranches.get(0);
            context = compileExpression(conditionalBranch.condition(), context);

            context = context.enterBlock(conditionalBranch.body().boundVariables(), conditionalBranch.body().source());
            context = compileBlockWithYield(conditionalBranch.body(), context);

            var ifTrueResult = context.exitBlock();
            var ifTrue = ifTrueResult.value;
            context = ifTrueResult.context;

            context = context.enterBlock(List.of(), conditionalBranch.source());
            context = compileIf(
                conditionalBranches.subList(1, conditionalBranches.size()),
                elseBody,
                context
            );

            var ifFalseResult = context.exitBlock();
            var ifFalse = ifFalseResult.value;
            context = ifFalseResult.context;

            context = context.addInstruction(new ShidIf(
                1,
                ifTrue,
                ifFalse,
                conditionalBranch.source()
            ));
        }

        return context;
    }

    private static Context compileIntLiteral(TypedIntLiteralNode node, Context context) {
        return context.addInstruction(new ShidPushInt(node.value(), node.source()));
    }

    private static Context compileMemberAccess(
        TypedMemberAccessNode node,
        Context context
    ) {
        context = compileExpression(node.receiver(), context);
        // TODO: handle not a type with member

        return switch (node.memberKind()) {
            case FIELD -> {
                var receiverType = node.receiver().type();
                var receiver = fieldReceiverTypeToReceiver(receiverType);

                yield context.addInstruction(new ShidFieldLoad(
                    receiver,
                    node.memberName(),
                    node.source()
                ));
            }
            case METHOD -> {
                // TODO: handle not a concrete type
                var structType = (SimpleConcreteType) Types.underlyingType(node.receiver().type());
                // TODO: handle not a function type
                var methodType = (FunctionType) node.type();

                var receiverLocalNameResult = context.generateShidVal("receiver");
                var receiverLocalName = receiverLocalNameResult.value();
                context = receiverLocalNameResult.context();

                context = context.addInstruction(
                    new ShidLocalStore(receiverLocalName, node.source())
                );

                var shedPositionalParamNames = new ArrayList<String>();
                for (
                    var positionalParamIndex = 0;
                    positionalParamIndex < methodType.positionalParams().size();
                    positionalParamIndex++
                ) {
                    shedPositionalParamNames.add("arg" + positionalParamIndex);
                }

                var boundVariables = Stream.concat(
                    shedPositionalParamNames.stream(),
                    methodType.namedParams().stream()
                        .map(namedParam -> namedParam.name())
                )
                    .map(BoundVariable::param)
                    .toList();

                context = context.enterBlock(boundVariables, node.source());

                context = context.addInstruction(
                    new ShidMethodLoad(structType, node.memberName(), node.source())
                );

                context = context.addInstruction(
                    new ShidLocalLoad(receiverLocalName, node.source())
                );

                var shidPositionalParams = new ArrayList<ShidParam>();
                for (var positionalParamName : shedPositionalParamNames) {
                    var shidParamName = context.shedVariableToLocal(positionalParamName);
                    shidPositionalParams.add(new ShidParam(shidParamName, node.source()));
                    context = context.addInstruction(
                        new ShidLocalLoad(shidParamName, node.source())
                    );
                }

                var shidNamedParams = new ArrayList<ShidParam>();
                for (var namedParam : methodType.namedParams()) {
                    var shidParamName = context.shedVariableToLocal(namedParam.name());
                    shidNamedParams.add(new ShidParam(shidParamName, node.source()));
                    context = context.addInstruction(
                        new ShidLocalLoad(shidParamName, node.source())
                    );
                }

                context = context.addInstruction(
                    new ShidCall(
                        methodType.positionalParams().size() + 1,
                        methodType.namedParams().stream()
                            .map(namedParam -> namedParam.name())
                            .toList(),
                        node.source()
                    )
                );

                context = context.addInstruction(new ShidReturn(node.source()));

                var bodyResult = context.exitBlock();
                var body = bodyResult.value();
                context = bodyResult.context();

                yield context.addInstruction(
                    new ShidDefineFunction(
                        node.memberName(),
                        shidPositionalParams,
                        shidNamedParams,
                        FreeVariables.builder()
                            .plus(receiverLocalName, node.source())
                            .build(),
                        body,
                        ShidDefineFunction.Variant.FUNCTION,
                        node.source()
                    )
                );
            }
        };
    }

    private static ShidFieldLoad.Receiver fieldReceiverTypeToReceiver(Type receiverType) {
        return switch (receiverType) {
            case ConstructedType<?> constructedType ->
                fieldReceiverTypeToReceiver(constructedType.genericType());

            case NamespaceType namespaceType ->
                new ShidFieldLoad.Receiver.Namespace(namespaceType.namespaceName());

            case SimpleStructType structType ->
                new ShidFieldLoad.Receiver.Struct(structType);

            case TypeLevelValueType typeLevelValueType -> {
                if (typeLevelValueType.value() instanceof UserDefinedEffect effect) {
                    yield new ShidFieldLoad.Receiver.Effect(effect);
                } else {
                    throw new RuntimeException("not a type with fields");
                }
            }

            default -> throw new RuntimeException("not a type with fields");
        };
    }

    public static ShidNamespace compileNamespace(TypedNamespaceNode node) {
        var context = Context.EMPTY.enterBlock(
            node.namespaceName(),
            node.boundVariables(),
            node.source()
        );

        context = context.addShedVal(ShedBuiltins.PRINT_NAME);
        var printLocalName = context.shedVariableToLocal(ShedBuiltins.PRINT_NAME);

        context = context
            .addInstruction(new ShidNamespaceImport(ShedBuiltins.NAMESPACE_NAME, node.source()))
            .addInstruction(new ShidFieldLoad(
                new ShidFieldLoad.Receiver.Namespace(ShedBuiltins.NAMESPACE_NAME),
                ShedBuiltins.PRINT_NAME,
                node.source()
            ))
            .addInstruction(new ShidLocalStore(printLocalName, node.source()));

        for (var importNode : node.imports()) {
            var localName = context.shedVariableToLocal(importNode.variableName());
            context = context
                .addInstruction(new ShidNamespaceImport(
                    importNode.namespaceName(),
                    importNode.source()
                ))
                .addInstruction(new ShidLocalStore(
                    localName,
                    importNode.source()
                ));

        }

        var statementsAfterHoisting = Iterables.sortedBy(node.body(), statement -> switch (statement) {
            case TypedAssignNode assign ->
                2;

            case TypedEffectDefinitionNode effectDefinition ->
                0;

            case TypedFunctionDefinitionNode functionDefinition ->
                1;

            case TypedNativeFunctionDeclarationNode nativeFunctionDeclaration ->
                1;

            case TypedNativeTypeDeclarationNode nativeTypeDeclaration ->
                0;

            case TypedStructDefinitionNode structDefinition ->
                0;

            case TypedSumTypeDefinitionNode sumTypeDefinition ->
                0;
        });
        for (var statement : statementsAfterHoisting) {
            context = compileNamespaceStatement(statement, context);
        }
        for (var statement : statementsAfterHoisting) {
            if (statement instanceof TypedStructDefinitionNode structDefinitionNode) {
                // TODO: test method hoisting
                context = compileStructDefinitionMethods(structDefinitionNode, context);
            } else if (statement instanceof TypedNativeTypeDeclarationNode nativeTypeDeclarationNode) {
                context = compileNativeTypeDeclarationMethods(nativeTypeDeclarationNode, context);
            }
        }

        for (var fieldName : node.fieldNames()) {
            var localName = context.shedVariableToLocal(fieldName);

            context = context
                .addInstruction(new ShidLocalLoad(localName, node.source()))
                .addInstruction(new ShidNamespaceFieldStore(
                    node.namespaceName(),
                    fieldName,
                    node.source()
                ));
        }

        var bodyResult = context.exitBlock();
        var body = bodyResult.value;
        context = bodyResult.context;

        return new ShidNamespace(node.namespaceName(), body, node.source());
    }

    public static Context compileNamespaceStatement(TypedNamespaceStatementNode node, Context context) {
        return switch (node) {
            case TypedAssignNode assign ->
                compileAssign(assign, context);

            case TypedEffectDefinitionNode effectDefinition ->
                compileEffectDefinition(effectDefinition, context);

            case TypedFunctionDefinitionNode functionDefinition ->
                compileFunctionDefinition(functionDefinition, context);

            case TypedNativeFunctionDeclarationNode nativeFunctionDeclaration ->
                compileNativeFunctionDeclaration(nativeFunctionDeclaration, context);

            case TypedNativeTypeDeclarationNode nativeTypeDeclarationNode ->
                compileNativeTypeDeclaration(nativeTypeDeclarationNode, context);

            case TypedStructDefinitionNode structDefinition ->
                compileStructDefinition(structDefinition, context);

            case TypedSumTypeDefinitionNode sumDefinition ->
                compileSumTypeDefinition(sumDefinition, context);
        };
    }

    private static Context compileNativeFunctionDeclaration(
        TypedNativeFunctionDeclarationNode node,
        Context context
    ) {
        var signature = node.signature();
        context = compileNativeFunction(
            context.namespaceName(),
            node.name(),
            signature,
            node.source(),
            context
        );

        var localName = context.shedVariableToLocal(node.name());
        var storeInstruction = new ShidLocalStore(localName, node.source());

        return context.addInstruction(storeInstruction);
    }

    private static Context compileNativeFunction(
        NamespaceName namespaceName,
        String name,
        TypedFunctionTypeNode signature,
        Source source,
        Context context
    ) {
        return context.addInstruction(new ShidDeclareNativeFunction(
            namespaceName,
            name,
            IntStream.range(0, signature.value().positionalParams().size())
                .mapToObj(paramIndex -> new ShidParam("arg" + paramIndex, signature.source()))
                .toList(),
            signature.value().namedParams().stream()
                .map(param -> new ShidParam(param.name(), signature.source()))
                .toList(),
            source
        ));
    }

    private static Context compileNativeTypeDeclaration(
        TypedNativeTypeDeclarationNode node,
        Context context
    ) {
        var nativeType = node.typeOrConstructor().toType();

        context = context.addInstruction(new ShidDeclareNativeType(
            nativeType,
            node.source()
        ));

        var localName = context.shedVariableToLocal(nativeType.name());
        var storeInstruction = new ShidLocalStore(localName, node.source());

        return context.addInstruction(storeInstruction);
    }

    static Context compileNativeTypeDeclarationMethods(
        TypedNativeTypeDeclarationNode node,
        Context context
    ) {
        var type = node.typeOrConstructor().toType();

        for (var methodNode : node.methods()) {
            context = compileNativeFunction(
                type.namespaceName(),
                String.format(
                    "method__%s__%s",
                    type.name(),
                    methodNode.name()
                ),
                methodNode.signature(),
                methodNode.source(),
                context
            );

            context = context.addInstruction(
                new ShidMethodStore(
                    type,
                    methodNode.name(),
                    methodNode.source()
                )
            );
        }

        return context;
    }

    private static Context compileReference(TypedReferenceNode node, Context context) {
        var shidVariable = context.shedVariableToShidVariable(node.name());

        if (shidVariable.isVar()) {
            return context.addInstruction(new ShidVarLoad(
                shidVariable.name(),
                node.source()
            ));
        } else {
            return context.addInstruction(new ShidLocalLoad(
                shidVariable.name(),
                node.source()
            ));
        }
    }

    private static Context compileResume(TypedResumeNode node, Context context) {
        context = compileExpression(node.expression(), context);
        context = context.addInstruction(new ShidResume(node.source()));
        return context;
    }

    private static Context compileStringLiteral(
        TypedStringLiteralNode node,
        Context context
    ) {
        var instruction = new ShidPushString(node.value(), node.source());
        return context.addInstruction(instruction);
    }

    private static Context compileStructDefinition(
        TypedStructDefinitionNode node,
        Context context
    ) {
        var structType = node.typeOrConstructor().toType();

        var defineStructInstruction = new ShidDefineStruct(
            structType,
            node.fields().stream()
                .map(field -> new ShidDefineStruct.Field(field.name(), field.source()))
                .toList(),
            node.source()
        );

        var localName = context.shedVariableToLocal(structType.name());
        var storeInstruction = new ShidLocalStore(localName, node.source());

        return context
            .addInstruction(defineStructInstruction)
            .addInstruction(storeInstruction);
    }

    static Context compileStructDefinitionMethods(
        TypedStructDefinitionNode node,
        Context context
    ) {
        var structType = node.typeOrConstructor().toType();

        for (var methodNode : node.methods()) {
            var shidFunctionName = methodToShidLocalName(structType, methodNode.name());
            context = compileFunction(
                shidFunctionName,
                methodNode.signature().params(),
                methodNode.body(),
                ShidDefineFunction.Variant.FUNCTION,
                methodNode.source(),
                context
            );

            context = context.addInstruction(
                new ShidMethodStore(
                    structType,
                    methodNode.name(),
                    methodNode.source()
                )
            );
        }

        return context;
    }

    private static String methodToShidLocalName(
        SimpleStructType structType,
        String methodName
    ) {
        return MessageFormat.format(
            "$methods__{0}__{1}__{2}",
            String.join("_", structType.namespaceName().parts()),
            structType.name(),
            methodName
        );
    }

    private static Context compileSumTypeDefinition(
        TypedSumTypeDefinitionNode node,
        Context context
    ) {
        var defineSumTypeInstruction = new ShidPushUnit(node.source());

        var sumType = node.typeOrConstructor().toType();
        var localName = context.shedVariableToLocal(sumType.name());
        var storeInstruction = new ShidLocalStore(localName, node.source());

        return context
            .addInstruction(defineSumTypeInstruction)
            .addInstruction(storeInstruction);
    }

    private static Context compileTarget(
        TypedTargetNode target,
        Context context
    ) {
        return switch (target) {
            case TypedTargetNode.Ignore ignore -> {
                yield context.addInstruction(new ShidDiscard(ignore.source()));
            }

            case TypedTargetNode.Reassign reassign -> {
                context = context.addInstruction(
                    new ShidVarStore(
                        context.shedVariableToLocal(reassign.name()),
                        reassign.source()
                    )
                );

                yield context;
            }

            case TypedTargetNode.Tuple tuple -> {
                if (tuple.elements().size() == 0) {
                    context = context.addInstruction(new ShidDiscard(tuple.source()));
                }

                for (var elementIndex = 0; elementIndex < tuple.elements().size(); elementIndex++) {
                    if (elementIndex < tuple.elements().size() - 1) {
                        context = context
                            .addInstruction(new ShidDuplicate(tuple.source()));
                    }

                    context = context
                        .addInstruction(new ShidTupleGet(
                            elementIndex,
                            tuple.elements().size(),
                            tuple.source()
                        ));

                    context = compileTarget(tuple.elements().get(elementIndex), context);
                }

                yield context;
            }

            case TypedTargetNode.Val val -> {
                context = context.addInstruction(
                    new ShidLocalStore(
                        context.shedVariableToLocal(val.name()),
                        val.source()
                    )
                );

                yield context;
            }

            case TypedTargetNode.Var var -> {
                context = context.addInstruction(
                    new ShidVarStore(
                        context.shedVariableToLocal(var.name()),
                        var.source()
                    )
                );

                yield context;
            }
        };
    }

    private static Context compileTry(TypedTryNode node, Context context) {
        for (var handlerMethod : node.handlerMethods()) {
            context = compileTryHandlerMethod(handlerMethod, context);
        }

        context = context.enterBlock(node.body().boundVariables(), node.body().source());
        context = compileBlockWithYield(node.body(), context);

        var bodyResult = context.exitBlock();
        var body = bodyResult.value;
        context = bodyResult.context;

        var instruction = new ShidTry(
            node.handledEffect().value(),
            node.handlerMethods().stream()
                .map(methodNode -> methodNode.name())
                .toList(),
            body,
            node.source()
        );

        return context.addInstruction(instruction);
    }

    private static Context compileTryHandlerMethod(
        TypedTryHandlerMethodNode node,
        Context context
    ) {
        return compileFunction(
            node.name(),
            node.params(),
            node.body(),
            ShidDefineFunction.Variant.EFFECT_HANDLER_METHOD,
            node.source(),
            context
        );
    }

    private static Context compileTuple(TypedTupleNode node, Context context) {
        for (var elementNode : node.elements()) {
            context = compileExpression(elementNode, context);
        }

        var instruction = new ShidTupleCreate(node.elements().size(), node.source());
        return context.addInstruction(instruction);
    }

    private static Context compileUnitLiteral(
        TypedUnitLiteralNode node,
        Context context
    ) {
        return context.addInstruction(new ShidPushUnit(node.source()));
    }

    private static Context compileWhen(TypedWhenNode node, Context context) {
        context = compileExpression(node.expression(), context);

        var switchCases = new ArrayList<ShidSwitch.Case>();

        for (var whenCase : node.cases()) {
            context = context.enterBlock(whenCase.body().boundVariables(), whenCase.body().source());
            context = compileBlockWithYield(whenCase.body(), context);

            var caseBodyResult = context.exitBlock();
            var caseBody = caseBodyResult.value;
            context = caseBodyResult.context;

            var switchCase = new ShidSwitch.Case(
                whenCase.caseType().rawType(),
                caseBody,
                whenCase.source()
            );
            switchCases.add(switchCase);
        }

        context = context.addInstruction(new ShidSwitch(
            switchCases,
            node.source()
        ));

        return context;
    }

    private static Context compileWhile(TypedWhileNode node, Context context) {
        context = context.enterBlock(node.body().boundVariables(), node.body().source());
        context = compileBlockWithoutYield(node.body(), context);

        var ifTrueResult = context.exitBlock();
        var ifTrue = ifTrueResult.value;
        context = ifTrueResult.context;

        context = context.enterBlock(List.of(), node.source());

        context = compileExpression(node.condition(), context);
        context = context.addInstruction(new ShidIf(
            0,
            new ShidBlock(List.of(), List.of(), node.source()),
            new ShidBlock(
                List.of(),
                List.of(
                    new ShidLoopBreak(node.source())
                ),
                node.source()
            ),
            node.source()
        ));

        context = context.addInstruction(ifTrue);

        context = context.addInstruction(new ShidLoopContinue(node.source()));

        var bodyResult = context.exitBlock();
        var body = bodyResult.value;
        context = bodyResult.context;

        context = context.addInstruction(new ShidLoop(body, node.source()));

        return context;
    }

    private static class BlockContext {
        public static final BlockContext EMPTY = new BlockContext(
            Optional.empty(),
            P.vector(),
            P.map(),
            P.vector(),
            Sources.NULL
        );

        private final Optional<NamespaceName> namespaceName;
        private final PVector<ShidInstruction> instructions;
        private final PMap<String, ShidVariable> shedVariablesByName;
        private final PVector<ShidVariable> boundVariables;
        private final Source source;

        public BlockContext(
            Optional<NamespaceName> namespaceName,
            PVector<ShidInstruction> instructions,
            PMap<String, ShidVariable> shedVariablesByName,
            PVector<ShidVariable> boundVariables,
            Source source
        ) {
            this.namespaceName = namespaceName;
            this.instructions = instructions;
            this.shedVariablesByName = shedVariablesByName;
            this.boundVariables = boundVariables;
            this.source = source;
        }

        public ShidBlock toBlock() {
            // TODO: test collection of bound locals
            return new ShidBlock(this.boundVariables, this.instructions, this.source);
        }

        public BlockContext addInstruction(ShidInstruction instruction) {
            return new BlockContext(
                this.namespaceName,
                this.instructions.plus(instruction),
                this.shedVariablesByName,
                this.boundVariables,
                this.source
            );
        }

        public ShidVariable shedVariable(String variableName) {
            var shedVariable = this.shedVariablesByName.get(variableName);

            if (shedVariable == null) {
                throw new RuntimeException("Unknown Shed variable: " + variableName);
            }

            return shedVariable;
        }

        public BlockContext addShedVariable(String shedName, ShidVariable shidVariable) {
            return new BlockContext(
                this.namespaceName,
                this.instructions,
                this.shedVariablesByName.plus(shedName, shidVariable),
                this.boundVariables.plus(shidVariable),
                this.source
            );
        }

        public BlockContext addShidVal(String localName) {
            return new BlockContext(
                this.namespaceName,
                this.instructions,
                this.shedVariablesByName,
                this.boundVariables.plus(new ShidVariable(localName, ShidVariable.Variant.VAL)),
                this.source
            );
        }
    }

    public static class Context {
        public static final Context EMPTY = new Context(
            P.stack(),
            // TODO: no block context by default?
            BlockContext.EMPTY,
            Empty.set()
        );

        private final PStack<BlockContext> blockStack;
        private final BlockContext block;
        private final PSet<String> localNames;

        private Context(
            PStack<BlockContext> blockStack,
            BlockContext block,
            PSet<String> localNames
        ) {
            this.blockStack = blockStack;
            this.block = block;
            this.localNames = localNames;
        }

        public Context enterBlock(
            List<BoundVariable> shedVariables,
            Source source
        ) {
            return enterBlock(Optional.empty(), shedVariables, source);
        }

        public Context enterBlock(
            NamespaceName namespaceName,
            List<BoundVariable> shedVariables,
            Source source
        ) {
            return enterBlock(
                Optional.of(namespaceName),
                shedVariables,
                source
            );
        }

        public Context enterBlock(
            Optional<NamespaceName> namespaceName,
            List<BoundVariable> shedVariables,
            Source source
        ) {
            var context = new Context(
                this.blockStack.plus(this.block),
                new BlockContext(
                    namespaceName,
                    P.vector(),
                    this.block.shedVariablesByName,
                    P.vector(),
                    source
                ),
                this.localNames
            );

            for (var shedVariable : shedVariables) {
                context = context.addShedVariable(shedVariable);
            }

            return context;
        }

        public Result<ShidBlock> exitBlock() {
            return new Context(
                this.blockStack.minus(0),
                this.blockStack.get(0),
                this.localNames
            ).toResult(this.block.toBlock());
        }

        private Context addInstruction(ShidInstruction instruction) {
            return new Context(
                this.blockStack,
                this.block.addInstruction(instruction),
                this.localNames
            );
        }

        public String shedVariableToLocal(String variableName) {
            return shedVariableToShidVariable(variableName).name();
        }

        public ShidVariable shedVariableToShidVariable(String variableName) {
            return this.block.shedVariable(variableName);
        }

        // TODO: make private / remove? Variables should be declared using enterBlock().
        public Context addShedVal(String variableName) {
            return addShedVariable(BoundVariable.val(variableName));
        }

        // TODO: make private? Variables should be declared using enterBlock().
        public Context addShedVariable(BoundVariable boundVariable) {
            var localName = generateLocalNameForVariable(boundVariable.name());

            return new Context(
                this.blockStack,
                this.block.addShedVariable(
                    boundVariable.name(),
                    new ShidVariable(localName, switch (boundVariable.variant()) {
                        case PARAM -> ShidVariable.Variant.PARAM;
                        case VAL -> ShidVariable.Variant.VAL;
                        case VAR -> ShidVariable.Variant.VAR;
                    })
                ),
                this.localNames.plus(localName)
            );
        }

        public Result<String> generateShidVal(String baseName) {
            var localName = generateLocalNameForVariable(baseName);

            return new Context(
                this.blockStack,
                this.block.addShidVal(localName),
                this.localNames.plus(localName)
            ).toResult(localName);
        }

        private String generateLocalNameForVariable(String variableName) {
            var index = 0;
            while (true) {
                var localName = variableName + "__" + index;
                if (this.localNames.contains(localName)) {
                    index++;
                } else {
                    return localName;
                }
            }
        }

        public List<ShidInstruction> instructions() {
            return this.block.instructions;
        }

        public <T> Result<T> toResult(T value) {
            return new Result<>(value, this);
        }

        public NamespaceName namespaceName() {
            // TODO: better error
            // TODO: walk up the block stack? Or make it clear that this will
            // only get the namespace for the immediate block.
            return this.block.namespaceName.orElseThrow();
        }

        public record Result<T>(T value, Context context) {}
    }

}