package org.shedlang.compiler.shid;

import org.shedlang.compiler.shid.instructions.*;
import org.shedlang.compiler.treetext.Tree;
import org.shedlang.compiler.types.*;

import java.util.List;

import static org.shedlang.compiler.treetext.TreeSerializer.serializeTree;
import static org.shedlang.compiler.treetext.Trees.*;

public class ShidSerializer {
    private ShidSerializer() {
    }

    public static void serializeBlock(ShidBlock block, StringBuilder builder) {
        var tree = blockToTree(block);
        serializeTree(tree, builder, 0);
    }

    public static void serializeInstruction(ShidInstruction instruction, StringBuilder builder) {
        var tree = instructionToTree(instruction);
        serializeTree(tree, builder, 0);
    }

    private static Tree blockToTree(ShidBlock block) {
        return treeListIndented(
            treeIdentifier("block"),
            treeListIndented(
                block.boundVariables().stream()
                    .map(variable -> treeList(
                        treeIdentifier(switch (variable.variant()) {
                            case PARAM -> "param";
                            case VAL -> "val";
                            case VAR -> "var";
                        }),
                        treeString(variable.name())
                    ))
                    .toList()
            ),
            instructionsToTree(block.instructions())
        );
    }

    private static Tree instructionsToTree(List<? extends ShidInstruction> instructions) {
        return treeListIndented(
            instructions.stream()
                .map(instruction -> instructionToTree(instruction))
                .toList()
        );
    }

    private static Tree instructionToTree(ShidInstruction instruction) {
        return switch (instruction) {
            case ShidBlock blockInstruction ->
                blockToTree(blockInstruction);

            case ShidCall callInstruction ->
                treeList(
                    treeIdentifier("call"),
                    treeNamed(
                        "positional-arg-count",
                        treeInt(callInstruction.positionalArgCount())
                    ),
                    treeNamed(
                        "named-arg-names",
                        treeList(
                            callInstruction.namedArgNames().stream()
                                .map(arg -> treeString(arg))
                                .toList()
                        )
                    )
                );

            case ShidDeclareNativeFunction declareNativeFunction ->
                treeListIndented(
                    treeIdentifier("declare-native-function"),
                    treeNamed("namespace-name", namespaceNameToTree(declareNativeFunction.namespaceName())),
                    treeNamed("name", treeString(declareNativeFunction.name())),
                    positionalParamsToTree(declareNativeFunction.positionalParams()),
                    namedParamsToTree(declareNativeFunction.namedParams())
                );

            case ShidDeclareNativeType declareNativeType ->
                treeList(
                    treeIdentifier("declare-native-type"),
                    nativeTypeToTree(declareNativeType.type())
                );

            case ShidDefineEffect defineEffectInstruction ->
                treeList(
                    treeIdentifier("define-effect"),
                    userDefinedEffectToTree(defineEffectInstruction.effect())
                );

            case ShidDefineFunction defineFunctionInstruction ->
                treeListIndented(
                    treeIdentifier("define-function"),
                    treeNamed("variant", switch (defineFunctionInstruction.variant()) {
                        case EFFECT_HANDLER_METHOD -> treeString("effect-handler-method");
                        case FUNCTION -> treeString("function");
                    }),
                    treeNamed("name", treeString(defineFunctionInstruction.name())),
                    positionalParamsToTree(defineFunctionInstruction.positionalParams()),
                    namedParamsToTree(defineFunctionInstruction.namedParams()),
                    treeNamed("free-variables", treeList(
                        defineFunctionInstruction.freeVariables().stream()
                            .map(freeVariable -> treeString(freeVariable.name()))
                            .toList()
                    )),
                    treeNamed("body", blockToTree(defineFunctionInstruction.body()))
                );

            case ShidDefineStruct defineStructInstruction ->
                treeListIndented(
                    treeIdentifier("define-struct"),
                    treeNamed("type", concreteTypeToTree(defineStructInstruction.type())),
                    treeNamed("fields", treeList(
                        defineStructInstruction.fields().stream()
                            .map(field -> treeString(field.name()))
                            .toList()
                    ))
                );

            case ShidDiscard discardInstruction ->
                treeList(
                    treeIdentifier("discard")
                );

            case ShidDuplicate duplicate ->
                treeList(
                    treeIdentifier("duplicate")
                );

            case ShidEffectHandlerBreak effectHandlerBreak ->
                treeList(
                    treeIdentifier("effect-handler-break"),
                    userDefinedEffectToTree(effectHandlerBreak.effect())
                );

            case ShidFieldLoad fieldLoad ->
                treeList(
                    treeIdentifier("field-load"),
                    switch (fieldLoad.receiver()) {
                        case ShidFieldLoad.Receiver.Effect effect ->
                            treeList(
                                treeIdentifier("effect"),
                                userDefinedEffectToTree(effect.effect())
                            );

                        case ShidFieldLoad.Receiver.Namespace namespace ->
                            treeList(
                                treeIdentifier("namespace"),
                                namespaceNameToTree(namespace.name())
                            );

                        case ShidFieldLoad.Receiver.Struct struct ->
                            treeList(
                                treeIdentifier("struct"),
                                concreteTypeToTree(struct.type())
                            );
                    },
                    treeString(fieldLoad.fieldName())
                );

            case ShidIf if_ ->
                treeListIndented(
                    treeIdentifier("if"),
                    treeNamed("result-count", treeInt(if_.resultCount())),
                    blockToTree(if_.trueBody()),
                    blockToTree(if_.falseBody())
                );

            case ShidIntAdd intAdd ->
                treeList(treeIdentifier("int-add"));

            case ShidIntEquals intEquals ->
                treeList(treeIdentifier("int-eq"));

            case ShidIntGreaterThan intGreaterThan ->
                treeList(treeIdentifier("int-gt"));

            case ShidIntGreaterThanOrEqual intGreaterThanOrEqual ->
                treeList(treeIdentifier("int-gte"));

            case ShidIntLessThan intLessThan ->
                treeList(treeIdentifier("int-lt"));

            case ShidIntLessThanOrEqual intLessThanOrEqual ->
                treeList(treeIdentifier("int-lte"));

            case ShidIntMultiply intMultiply ->
                treeList(treeIdentifier("int-mul"));

            case ShidIntNotEqual intNotEqual ->
                treeList(treeIdentifier("int-ne"));

            case ShidIntSub intSub ->
                treeList(treeIdentifier("int-sub"));

            case ShidLocalDrop localDrop ->
                treeList(
                    treeIdentifier("local-drop"),
                    treeString(localDrop.name())
                );

            case ShidLocalLoad localLoadInstruction ->
                treeList(
                    treeIdentifier("local-load"),
                    treeString(localLoadInstruction.name())
                );

            case ShidLocalStore localStoreInstruction ->
                treeList(
                    treeIdentifier("local-store"),
                    treeString(localStoreInstruction.name())
                );

            case ShidLoop loop ->
                treeListIndented(
                    treeIdentifier("loop"),
                    blockToTree(loop.body())
                );

            case ShidLoopBreak loopBreak ->
                treeList(treeIdentifier("break"));

            case ShidLoopContinue loopContinue ->
                treeList(treeIdentifier("continue"));

            case ShidMethodLoad methodLoad ->
                treeList(
                    treeIdentifier("method-load"),
                    concreteTypeToTree(methodLoad.type()),
                    treeString(methodLoad.methodName())
                );

            case ShidMethodStore methodStore ->
                treeList(
                    treeIdentifier("method-store"),
                    concreteTypeToTree(methodStore.type()),
                    treeString(methodStore.methodName())
                );

            case ShidNamespaceFieldStore namespaceFieldStore ->
                treeList(
                    treeIdentifier("namespace-field-store"),
                    namespaceNameToTree(namespaceFieldStore.namespaceName()),
                    treeString(namespaceFieldStore.fieldName())
                );

            case ShidNamespaceImport namespaceImport ->
                treeList(
                    treeIdentifier("namespace-import"),
                    namespaceNameToTree(namespaceImport.namespaceName())
                );

            case ShidPushBool pushBoolInstruction ->
                treeList(
                    treeIdentifier("push-bool"),
                    treeIdentifier(Boolean.toString(pushBoolInstruction.value()))
                );

            case ShidPushInt pushIntInstruction ->
                treeList(
                    treeIdentifier("push-int"),
                    treeIdentifier(pushIntInstruction.value().toString())
                );

            case ShidPushString pushStringInstruction ->
                treeList(
                    treeIdentifier("push-string"),
                    treeString(pushStringInstruction.value())
                );

            case ShidPushUnit pushUnitInstruction ->
                treeList(
                    treeIdentifier("push-unit")
                );

            case ShidResume resumeInstruction ->
                treeList(
                    treeIdentifier("resume")
                );

            case ShidReturn returnInstruction ->
                treeList(
                    treeIdentifier("return")
                );

            case ShidStringAdd stringAddInstruction ->
                treeList(treeIdentifier("string-add"));

            case ShidSwitch switchInstruction ->
                treeListIndented(
                    treeIdentifier("switch"),
                    treeListIndented(
                        switchInstruction.cases().stream()
                            .map(switchCase -> treeListIndented(
                                concreteTypeToTree(switchCase.caseType()),
                                blockToTree(switchCase.body())
                            ))
                            .toList()
                    )
                );

            case ShidTry tryInstruction ->
                treeListIndented(
                    treeIdentifier("try"),
                    userDefinedEffectToTree(tryInstruction.handledEffect()),
                    treeListIndented(
                        tryInstruction.handlerMethodNames().stream()
                            .map(name -> treeString(name))
                            .toList()
                    ),
                    blockToTree(tryInstruction.body())
                );

            case ShidTupleCreate tupleCreateInstruction ->
                treeList(
                    treeIdentifier("tuple-create"),
                    treeInt(tupleCreateInstruction.elementCount())
                );

            case ShidTupleGet tupleGetInstruction ->
                treeList(
                    treeIdentifier("tuple-get"),
                    treeInt(tupleGetInstruction.elementIndex()),
                    treeInt(tupleGetInstruction.elementCount())
                );

            case ShidVarLoad varLoadInstruction ->
                treeList(
                    treeIdentifier("var-load"),
                    treeString(varLoadInstruction.name())
                );

            case ShidVarStore varStoreInstruction ->
                treeList(
                    treeIdentifier("var-store"),
                    treeString(varStoreInstruction.name())
                );
        };
    }

    private static Tree namespaceNameToTree(NamespaceName namespaceName) {
        return treeString(namespaceName.toString());
    }

    private static Tree positionalParamsToTree(List<ShidParam> positionalParams) {
        return treeNamed("positional-params", treeList(
            positionalParams.stream()
                .map(param -> treeString(param.name()))
                .toList()
        ));
    }

    private static Tree namedParamsToTree(List<ShidParam> namedParams) {
        return treeNamed("named-params", treeList(
            namedParams.stream()
                .map(param -> treeString(param.name()))
                .toList()
        ));
    }

    private static Tree userDefinedEffectToTree(UserDefinedEffect effect) {
        return treeString(effect.describe());
    }

    private static Tree nativeTypeToTree(SimpleNativeType type) {
        return treeString(type.describe());
    }

    private static Tree concreteTypeToTree(SimpleConcreteType type) {
        return treeString(type.describe());
    }
}