package org.shedlang.compiler.shidtopatio;

import org.shedlang.compiler.ast.typed.FreeVariables;
import org.shedlang.compiler.patio.PatioExceptionTypeDefinition;
import org.shedlang.compiler.patio.PatioFunctionDeclaration;
import org.shedlang.compiler.patio.PatioLocal;
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.PatioType;
import org.shedlang.compiler.patio.types.PatioTypes;
import org.shedlang.compiler.shid.instructions.ShidParam;
import org.shedlang.compiler.sources.Source;
import org.shedlang.compiler.sources.Sources;
import org.shedlang.compiler.types.UserDefinedEffect;
import org.shedlang.compiler.types.UserDefinedEffectInfo;
import org.shedlang.compiler.util.Lists;

import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.IntStream;

public class ShidToPatioEffects {
    private static final PatioType EFFECT_TAG_TYPE = PatioTypes.INT;
    private static final PatioType EFFECT_METHOD_INDEX_TYPE = PatioTypes.INT;
    private static final PatioType EFFECT_HANDLER_TYPE = PatioTypes.DATA_PTR;
    private static final PatioType EFFECT_HANDLER_METHOD_POINTER_TYPE = PatioTypes.DATA_PTR;
    private static final PatioType EFFECT_HANDLER_METHOD_TYPE = ShidToPatioTypes.VALUE_TYPE;

    private static ShidToPatioClosures.ClosureRuntimeType effectHandlerMethodRuntimeType(
        FreeVariables freeVariables,
        ShidToPatioContext context
    ) {
        return ShidToPatioClosures.closureRuntimeType(
            EFFECT_HANDLER_METHOD_ADDITIONAL_DATA_LAYOUT,
            freeVariables,
            context
        );
    }

    private static final PatioDataLayout EFFECT_HANDLER_METHOD_ADDITIONAL_DATA_LAYOUT =
        PatioDataLayouts.struct(List.of(PatioDataLayouts.DATA_PTR));

    private static final List<Integer> CLOSURE_EFFECT_HANDLER_OFFSET
        = ShidToPatioClosures.additionalOffset(List.of(0));

    private static UserDefinedEffectRuntimeType userDefinedEffectRuntimeType(
        UserDefinedEffect effect,
        ShidToPatioContext context
    ) {
        var effectInfo = context.userDefinedEffectInfo(effect);
        return new UserDefinedEffectRuntimeType(effectInfo.methodCount());
    }

    private record UserDefinedEffectRuntimeType(int methodCount) implements RuntimeType {
        @Override
        public PatioDataLayout layout() {
            return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(Lists.repeat(
                PatioDataLayouts.fromType(ShidToPatioClosures.CLOSURE_TYPE),
                methodCount
            )));
        }

        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 methodCount;
        }

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

    public static EffectHandlerRuntimeType effectHandlerRuntimeType(
        UserDefinedEffect effect,
        UserDefinedEffectInfo effectInfo
    ) {
        return new EffectHandlerRuntimeType(effect, effectInfo.methodCount());
    }

    public record EffectHandlerRuntimeType(
        UserDefinedEffect effect,
        int methodCount
    ) implements RuntimeType {
        @Override
        public PatioDataLayout layout() {
            return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(List.of(
                PatioDataLayouts.fromType(EFFECT_HANDLER_TYPE),
                PatioDataLayouts.struct(Lists.repeat(
                    PatioDataLayouts.fromType(EFFECT_HANDLER_METHOD_TYPE),
                    methodCount
                ))
            )));
        }

        private static final List<Integer> GC_CHILDREN_OFFSET =
            ShidToPatioGc.gcObjectOffset(List.of(1));

        @Override
        public List<Integer> gcChildrenOffset() {
            return GC_CHILDREN_OFFSET;
        }

        @Override
        public int gcChildrenCount() {
            return methodCount;
        }

        public List<Integer> handlerMethodOffset(int handlerMethodIndex) {
            return ShidToPatioGc.gcObjectOffset(List.of(1, handlerMethodIndex));
        }
    }

    private static final PatioFunctionDeclaration FIND_HANDLER_METHOD_DECLARATION = new PatioFunctionDeclaration(
        "shed__runtime__effects_find_handler_method",
        List.of(
            PatioLocal.param("effectTag", EFFECT_TAG_TYPE),
            PatioLocal.param("methodIndex", EFFECT_METHOD_INDEX_TYPE)
        ),
        List.of(EFFECT_HANDLER_METHOD_TYPE),
        Sources.RUNTIME
    );

    private static final PatioFunctionDeclaration GET_HANDLER_DECLARATION = new PatioFunctionDeclaration(
        "shed__runtime__effects_get_handler",
        List.of(),
        List.of(EFFECT_HANDLER_TYPE),
        Sources.RUNTIME
    );

    private static final PatioFunctionDeclaration SET_HANDLER_DECLARATION = new PatioFunctionDeclaration(
        "shed__runtime__effects_set_handler",
        List.of(PatioLocal.param("effectHandler", EFFECT_HANDLER_TYPE)),
        List.of(),
        Sources.RUNTIME
    );

    private static final PatioFunctionDeclaration PUSH_HANDLER_DECLARATION = new PatioFunctionDeclaration(
        "shed__runtime__effects_push_handler",
        List.of(
            PatioLocal.param("effectHandler", EFFECT_HANDLER_TYPE)
        ),
        List.of(),
        Sources.RUNTIME
    );

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

    private ShidToPatioEffects() {
    }

    public static ShidToPatioContext createUserDefinedEffect(
        PatioLocalRef destination,
        UserDefinedEffect effect,
        Source source,
        ShidToPatioContext context
    ) {
        context = context.addCompilationUnitElement(new PatioExceptionTypeDefinition(
            context.exceptionTypeName(effect),
            ShidToPatioTypes.VALUE_TYPE,
            source
        ));

        var effectPointerResult = context.allocateVal(
            ShidToPatioLocalSpace.USER_DEFINED_EFFECT_POINTER,
            effect.name(),
            PatioTypes.DATA_PTR
        );
        var effectPointer = effectPointerResult.value();
        context = effectPointerResult.context();

        var effectRuntimeType = userDefinedEffectRuntimeType(effect, context);

        var effectInfo = context.userDefinedEffectInfo(effect);

        context = ShidToPatioGc.alloc(
            effectPointer,
            effectRuntimeType,
            source,
            context
        );

        for (var method : effectInfo.methods()) {
            var effectMethodResult = context.allocateVal(
                ShidToPatioLocalSpace.USER_DEFINED_EFFECT_METHOD,
                method.name(),
                ShidToPatioClosures.CLOSURE_TYPE
            );
            var effectMethod = effectMethodResult.value();
            context = effectMethodResult.context();

            var positionalParams = IntStream.range(0, method.type().positionalParams().size())
                // TODO: ensure names of positional params doesn't clash with named params
                .mapToObj(paramIndex -> new ShidParam("arg" + paramIndex, source))
                .toList();

            context = ShidToPatioFunctions.create(
                effectMethod,
                method.name(),
                positionalParams,
                method.type().namedParams().stream()
                    .map(namedParam -> new ShidParam(namedParam.name(), source))
                    .toList(),
                FreeVariables.empty(),
                bodyContext -> {
                    bodyContext = bodyContext.ensureCompilationUnitElement(FIND_HANDLER_METHOD_DECLARATION);

                    var handlerMethodResult = bodyContext.allocateVal(
                        ShidToPatioLocalSpace.EFFECT_HANDLER_METHOD,
                        EFFECT_HANDLER_METHOD_TYPE
                    );
                    var handlerMethod = handlerMethodResult.value();
                    bodyContext = handlerMethodResult.context();

                    bodyContext = bodyContext.addInstruction(new PatioCall(
                        List.of(handlerMethod),
                        FIND_HANDLER_METHOD_DECLARATION.name(),
                        List.of(
                            PatioIntConst.of(bodyContext.effectTag(effect)),
                            PatioIntConst.of(method.fieldIndex())
                        ),
                        source
                    ));

                    var returnValueResult = bodyContext.allocateVal(
                        ShidToPatioLocalSpace.TEMPORARY,
                        ShidToPatioTypes.VALUE_TYPE
                    );
                    var returnValue = returnValueResult.value();
                    bodyContext = returnValueResult.context();

                    var bodyContextArgs = bodyContext;
                    bodyContext = ShidToPatioClosures.call(
                        returnValue,
                        handlerMethod,
                        positionalParams.stream()
                            .map(param -> bodyContextArgs.shidLocalNameToLocal(param.name()))
                            .toList(),
                        method.type().namedParams().stream()
                            .map(param -> Map.entry(
                                param.name(),
                                bodyContextArgs.shidLocalNameToLocal(param.name())
                            ))
                            .toList(),
                        source,
                        bodyContext
                    );

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

                    return bodyContext;
                },
                source,
                context
            );

            context = context
                .addInstruction(new PatioMemoryStore(
                    effectPointer,
                    effectRuntimeType.layout(),
                    effectRuntimeType.methodOffset(method.fieldIndex()),
                    effectMethod,
                    source
                ));
        }

        context = context.addInstruction(new PatioDataPtrToInt(
            destination,
            effectPointer,
            source
        ));

        return context;
    }

    public static ShidToPatioContext fieldLoad(
        PatioLocalRef destination,
        PatioOperand receiver,
        UserDefinedEffect effect,
        String fieldName,
        Source source,
        ShidToPatioContext context
    ) {
        var effectPointerResult = context.allocateVal(
            ShidToPatioLocalSpace.USER_DEFINED_EFFECT_POINTER,
            PatioTypes.DATA_PTR
        );
        var effectPointer = effectPointerResult.value();
        context = effectPointerResult.context();

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

        var effectRuntimeType = userDefinedEffectRuntimeType(effect, context);

        var effectInfo = context.userDefinedEffectInfo(effect);
        var methodIndex = effectInfo.methodIndex(fieldName);

        context = context.addInstruction(new PatioMemoryLoad(
            destination,
            effectPointer,
            effectRuntimeType.layout(),
            effectRuntimeType.methodOffset(methodIndex),
            source
        ));

        return context;
    }

    public static ShidToPatioContext try_(
        PatioLocalRef destination,
        UserDefinedEffect effect,
        Map<String, PatioOperand> handlerMethodsByName,
        Function<ShidToPatioContext, ShidToPatioResult<PatioBlock>> generateBody,
        Source source,
        ShidToPatioContext context
    ) {
        var effectInfo = context.userDefinedEffectInfo(effect);

        context = context
            .ensureCompilationUnitElement(PUSH_HANDLER_DECLARATION)
            .ensureCompilationUnitElement(POP_HANDLER_DECLARATION);

        var effectHandlerResult = context.allocateVal(
            ShidToPatioLocalSpace.EFFECT_HANDLER,
            EFFECT_HANDLER_TYPE
        );
        var effectHandler = effectHandlerResult.value();
        context = effectHandlerResult.context();

        var effectHandlerRuntimeType = effectHandlerRuntimeType(effect, effectInfo);

        context = ShidToPatioGc.alloc(
            effectHandler,
            effectHandlerRuntimeType,
            source,
            context
        );

        context = context.addInstruction(new PatioCall(
            List.of(),
            PUSH_HANDLER_DECLARATION.name(),
            List.of(
                effectHandler
            ),
            source
        ));

        for (var handlerMethodInfo : effectInfo.methods()) {
            context = context.addInstruction(new PatioMemoryStore(
                effectHandler,
                effectHandlerRuntimeType.layout(),
                effectHandlerRuntimeType.handlerMethodOffset(handlerMethodInfo.fieldIndex()),
                handlerMethodsByName.get(handlerMethodInfo.name()),
                source
            ));
        }

        var popHandler = new PatioCall(
            List.of(),
            POP_HANDLER_DECLARATION.name(),
            List.of(),
            source
        );

        context = context.pushOnReturn(popHandler);

        var bodyResult = generateBody.apply(context);
        var body = bodyResult.value();
        context = bodyResult.context();

        context = context.popOnReturn();

        var exceptionValueResult = context.allocateVal(
            ShidToPatioLocalSpace.TEMPORARY,
            ShidToPatioTypes.VALUE_TYPE
        );
        var exceptionValue = exceptionValueResult.value();
        context = exceptionValueResult.context();

        context = context.addInstruction(new PatioTry(
            List.of(destination),
            List.of(
                new PatioCatchClause(
                    context.exceptionTypeName(effect),
                    exceptionValue,
                    new PatioBlock(
                        List.of(),
                        List.of(exceptionValue),
                        source
                    ),
                    source
                )
            ),
            body.addInstruction(popHandler),
            source
        ));

        return context;
    }

    public static ShidToPatioContext createEffectHandlerMethod(
        PatioLocalRef destination,
        String name,
        List<ShidParam> positionalParams,
        List<ShidParam> namedParams,
        FreeVariables freeVariables,
        Function<ShidToPatioContext, ShidToPatioContext> generateBody,
        Source source,
        ShidToPatioContext context
    ) {
        context = context
            .ensureCompilationUnitElement(GET_HANDLER_DECLARATION)
            .ensureCompilationUnitElement(SET_HANDLER_DECLARATION);

        var effectHandlerMethodRuntimeType = effectHandlerMethodRuntimeType(freeVariables, context);

        return ShidToPatioClosures.create(
            destination,
            ShidToPatioCompilationUnitElementSpace.EFFECT_HANDLER_METHOD,
            name,
            positionalParams,
            namedParams,
            effectHandlerMethodRuntimeType,
            (closurePointer, allocContext) -> {
                var effectHandlerResult = allocContext.allocateVal(
                    ShidToPatioLocalSpace.EFFECT_HANDLER_CURRENT,
                    EFFECT_HANDLER_TYPE
                );
                var effectHandler = effectHandlerResult.value();
                allocContext = effectHandlerResult.context();

                allocContext = getCurrentEffectHandler(effectHandler, source, allocContext);

                allocContext = allocContext.addInstruction(new PatioMemoryStore(
                    closurePointer,
                    effectHandlerMethodRuntimeType.layout(),
                    CLOSURE_EFFECT_HANDLER_OFFSET,
                    effectHandler,
                    source
                ));

                return allocContext;
            },
            freeVariables,
            (closurePointer, bodyContext) -> {
                var previousEffectHandlerResult = bodyContext.allocateVal(
                    ShidToPatioLocalSpace.EFFECT_HANDLER,
                    EFFECT_HANDLER_TYPE
                );
                var previousEffectHandler = previousEffectHandlerResult.value();
                bodyContext = previousEffectHandlerResult.context();

                bodyContext = getCurrentEffectHandler(previousEffectHandler, source, bodyContext);

                bodyContext = bodyContext.stashEffectHandler(previousEffectHandler);

                var effectHandlerResult = bodyContext.allocateVal(
                    ShidToPatioLocalSpace.EFFECT_HANDLER,
                    EFFECT_HANDLER_TYPE
                );
                var effectHandler = effectHandlerResult.value();
                bodyContext = effectHandlerResult.context();

                bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                    effectHandler,
                    closurePointer,
                    effectHandlerMethodRuntimeType.layout(),
                    CLOSURE_EFFECT_HANDLER_OFFSET,
                    source
                ));

                bodyContext = bodyContext.addInstruction(new PatioCall(
                    List.of(),
                    SET_HANDLER_DECLARATION.name(),
                    List.of(effectHandler),
                    source
                ));

                bodyContext = generateBody.apply(bodyContext);

                return bodyContext;
            },
            source,
            context
        );
    }

    public static ShidToPatioContext resume(
        PatioOperand value,
        Source source,
        ShidToPatioContext context
    ) {
        context = context.addInstruction(new PatioCall(
            List.of(),
            SET_HANDLER_DECLARATION.name(),
            List.of(context.stashedEffectHandler()),
            source
        ));

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

        return context;
    }

    public static ShidToPatioContext break_(
        UserDefinedEffect effect,
        PatioOperand value,
        Source source,
        ShidToPatioContext context
    ) {
        var exceptionTypeName = context.exceptionTypeName(effect);

        context = context.addInstruction(new PatioThrow(
            exceptionTypeName,
            value,
            source
        ));

        return context;
    }

    private static ShidToPatioContext getCurrentEffectHandler(
        PatioLocalRef destination,
        Source source,
        ShidToPatioContext context
    ) {
        return context.addInstruction(new PatioCall(
            List.of(destination),
            GET_HANDLER_DECLARATION.name(),
            List.of(),
            source
        ));
    }
}