package p.studio.compiler.pbs.lowering;

import p.studio.compiler.models.IRBackendExecutableFunction;
import p.studio.compiler.models.IRReservedMetadata;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsSemanticsErrors;
import p.studio.compiler.source.diagnostics.DiagnosticPhase;
import p.studio.compiler.source.identifiers.CallableId;
import p.studio.compiler.source.identifiers.NameId;
import p.studio.compiler.source.tables.NameTable;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static p.studio.compiler.pbs.lowering.PbsExecutableLoweringModels.*;

final class PbsExecutableCallsiteEmitter {
    CallableId resolveUniqueCallableId(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var calleeIdentity = resolveCalleeIdentity(callExpr.callee(), context.nameTable());
        if (calleeIdentity == null) {
            return null;
        }
        final var candidates = resolveCallsiteCandidates(callExpr, calleeIdentity, context);
        if (!candidates.hostCandidates().isEmpty() || !candidates.intrinsicCandidates().isEmpty()) {
            return null;
        }
        if (candidates.callableCandidates().size() != 1) {
            return null;
        }
        return candidates.callableCandidates().getFirst();
    }

    IRReservedMetadata.HostMethodBinding resolveUniqueHostBinding(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var calleeIdentity = resolveCalleeIdentity(callExpr.callee(), context.nameTable());
        if (calleeIdentity == null) {
            return null;
        }
        final var candidates = resolveCallsiteCandidates(callExpr, calleeIdentity, context);
        if (!candidates.callableCandidates().isEmpty() || !candidates.intrinsicCandidates().isEmpty()) {
            return null;
        }
        if (candidates.hostCandidates().size() != 1) {
            return null;
        }
        return candidates.hostCandidates().getFirst();
    }

    int returnSlotsOf(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var calleeIdentity = resolveCalleeIdentity(callExpr.callee(), context.nameTable());
        if (calleeIdentity == null) {
            return 0;
        }
        final var candidates = resolveCallsiteCandidates(callExpr, calleeIdentity, context);
        final var resolvedCategory = resolveCallsiteCategory(candidates, callExpr, context);
        if (resolvedCategory == null) {
            return 0;
        }
        return switch (resolvedCategory) {
            case HOST -> candidates.hostCandidates().getFirst().retSlots();
            case INTRINSIC -> candidates.intrinsicCandidates().getFirst().retSlots();
            case CALLABLE -> {
                if (candidates.callableCandidates().size() != 1) {
                    yield 0;
                }
                yield context.returnSlotsByCallableId().getOrDefault(candidates.callableCandidates().getFirst(), 0);
            }
        };
    }

    void emitCallsite(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var calleeIdentity = resolveCalleeIdentity(callExpr.callee(), context.nameTable());
        if (calleeIdentity == null) {
            reportUnsupportedLowering("executable lowering requires resolvable callee identity", callExpr.span(), context);
            return;
        }
        final var candidates = resolveCallsiteCandidates(callExpr, calleeIdentity, context);
        final var resolvedCategory = resolveCallsiteCategory(candidates, callExpr, context);
        if (resolvedCategory == null) {
            return;
        }
        switch (resolvedCategory) {
            case HOST -> emitHostCall(callExpr, context, candidates.hostCandidates().getFirst());
            case INTRINSIC -> emitIntrinsicCall(callExpr, context, candidates.intrinsicCandidates().getFirst());
            case CALLABLE -> emitCallableCall(callExpr, context, candidates.callableCandidates());
        }
    }

    private PbsResolvedCallsiteCandidates resolveCallsiteCandidates(
            final PbsAst.CallExpr callExpr,
            final PbsCalleeIdentity calleeIdentity,
            final PbsExecutableLoweringContext context) {
        return new PbsResolvedCallsiteCandidates(
                resolveCallableCandidates(callExpr, calleeIdentity, context),
                resolveHostCandidates(callExpr, calleeIdentity, context),
                resolveIntrinsicCandidates(callExpr, calleeIdentity, context));
    }

    private List<IRReservedMetadata.HostMethodBinding> resolveHostCandidates(
            final PbsAst.CallExpr callExpr,
            final PbsCalleeIdentity calleeIdentity,
            final PbsExecutableLoweringContext context) {
        final var hostCandidates = context.hostByMethodName().getOrDefault(calleeIdentity.memberNameId(), List.of());
        if (hostCandidates.isEmpty()) {
            return List.of();
        }
        final var ownerHint = hostOwnerHint(callExpr.callee());
        if (ownerHint.isBlank()) {
            return hostCandidates;
        }
        final var filtered = hostCandidates.stream()
                .filter(candidate -> ownerHint.equals(candidate.ownerName()))
                .toList();
        return filtered;
    }

    private List<CallableId> resolveCallableCandidates(
            final PbsAst.CallExpr callExpr,
            final PbsCalleeIdentity calleeIdentity,
            final PbsExecutableLoweringContext context) {
        final var callableCandidates = new ArrayList<CallableId>();
        appendCallableCandidates(
                callableCandidates,
                context.callableIdsByNameAndArity().get(
                        new PbsCallableResolutionKey(calleeIdentity.primaryCallableNameId(), callExpr.arguments().size())));
        if (calleeIdentity.secondaryCallableNameId() != null) {
            appendCallableCandidates(
                    callableCandidates,
                    context.callableIdsByNameAndArity().get(
                            new PbsCallableResolutionKey(calleeIdentity.secondaryCallableNameId(), callExpr.arguments().size())));
        }
        return callableCandidates;
    }

    private void appendCallableCandidates(
            final List<CallableId> destination,
            final List<CallableId> candidates) {
        if (candidates == null) {
            return;
        }
        for (final var candidate : candidates) {
            if (!destination.contains(candidate)) {
                destination.add(candidate);
            }
        }
    }

    private PbsCallsiteCategory resolveCallsiteCategory(
            final PbsResolvedCallsiteCandidates candidates,
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var categoryCount = countResolvedCategories(candidates);
        if (categoryCount == 0) {
            reportUnsupportedLowering("executable lowering requires resolvable callee identity", callExpr.span(), context);
            return null;
        }
        if (categoryCount > 1) {
            reportUnsupportedLowering("executable lowering found ambiguous callsite category", callExpr.span(), context);
            return null;
        }
        if (!candidates.hostCandidates().isEmpty()) {
            if (candidates.hostCandidates().size() > 1) {
                reportUnsupportedLowering("executable lowering found ambiguous host binding identity", callExpr.span(), context);
                return null;
            }
            return PbsCallsiteCategory.HOST;
        }
        if (!candidates.intrinsicCandidates().isEmpty()) {
            if (candidates.intrinsicCandidates().size() > 1) {
                reportUnsupportedLowering("executable lowering found ambiguous intrinsic identity", callExpr.span(), context);
                return null;
            }
            return PbsCallsiteCategory.INTRINSIC;
        }
        return PbsCallsiteCategory.CALLABLE;
    }

    private int countResolvedCategories(final PbsResolvedCallsiteCandidates candidates) {
        var categoryCount = 0;
        if (!candidates.callableCandidates().isEmpty()) {
            categoryCount++;
        }
        if (!candidates.hostCandidates().isEmpty()) {
            categoryCount++;
        }
        if (!candidates.intrinsicCandidates().isEmpty()) {
            categoryCount++;
        }
        return categoryCount;
    }

    private void emitHostCall(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context,
            final IRReservedMetadata.HostMethodBinding host) {
        final var effectiveArgSlots = callExpr.arguments().size() + implicitReceiverArgSlots(callExpr.callee(), context);
        context.instructions().add(new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.CALL_HOST,
                "",
                new IRBackendExecutableFunction.HostCallMetadata(
                        host.abiModule(),
                        host.abiMethod(),
                        host.abiVersion(),
                        effectiveArgSlots,
                        host.retSlots()),
                null,
                callExpr.span()));
    }

    private void emitIntrinsicCall(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context,
            final IRReservedMetadata.IntrinsicSurface intrinsic) {
        final var effectiveArgSlots = intrinsic.argSlots() + implicitReceiverArgSlots(callExpr.callee(), context);
        context.instructions().add(new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.CALL_INTRINSIC,
                "",
                null,
                new IRBackendExecutableFunction.IntrinsicCallMetadata(
                        intrinsic.canonicalName(),
                        intrinsic.canonicalVersion(),
                        context.intrinsicIdTable().register(intrinsic.canonicalName(), intrinsic.canonicalVersion())),
                effectiveArgSlots,
                intrinsic.retSlots(),
                callExpr.span()));
    }

    private void emitCallableCall(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context,
            final List<CallableId> callableCandidates) {
        if (callableCandidates.isEmpty()) {
            reportUnsupportedLowering("executable lowering requires resolvable callable identity", callExpr.span(), context);
            return;
        }
        if (callableCandidates.size() > 1) {
            reportUnsupportedLowering("executable lowering found ambiguous callable identity", callExpr.span(), context);
            return;
        }
        final var calleeCallableId = callableCandidates.getFirst();
        final var calleeSignature = context.callableSignatureByCallableId().get(calleeCallableId);
        if (calleeSignature == null) {
            reportUnsupportedLowering("executable lowering resolved callable without signature metadata", callExpr.span(), context);
            return;
        }
        final var effectiveArgSlots = callExpr.arguments().size() + implicitReceiverArgSlots(callExpr.callee(), context);
        context.instructions().add(new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.CALL_FUNC,
                calleeSignature.moduleId(),
                calleeSignature.callableName(),
                calleeCallableId,
                null,
                null,
                effectiveArgSlots,
                context.returnSlotsByCallableId().get(calleeCallableId),
                callExpr.span()));
    }

    private List<IRReservedMetadata.IntrinsicSurface> resolveIntrinsicCandidates(
            final PbsAst.CallExpr callExpr,
            final PbsCalleeIdentity calleeIdentity,
            final PbsExecutableLoweringContext context) {
        final var receiverOwner = resolveReceiverOwnerForCallee(callExpr.callee(), context);
        if (receiverOwner.isBlank()) {
            return List.of();
        }
        final var key = new PbsIntrinsicOwnerMethodKey(receiverOwner, calleeIdentity.memberSourceMethodName());
        return context.intrinsicByOwnerAndMethod().getOrDefault(key, List.of());
    }

    private String resolveReceiverOwnerForCallee(
            final PbsAst.Expression callee,
            final PbsExecutableLoweringContext context) {
        return switch (callee) {
            case PbsAst.MemberExpr memberExpr -> resolveExpressionOwner(memberExpr.receiver(), context);
            case PbsAst.GroupExpr groupExpr -> resolveReceiverOwnerForCallee(groupExpr.expression(), context);
            default -> "";
        };
    }

    String resolveExpressionOwnerForLowering(
            final PbsAst.Expression expression,
            final PbsExecutableLoweringContext context) {
        return resolveExpressionOwner(expression, context);
    }

    private String hostOwnerHint(final PbsAst.Expression callee) {
        return switch (callee) {
            case PbsAst.MemberExpr memberExpr -> {
                final var rootName = memberRootName(memberExpr.receiver());
                yield rootName == null ? "" : rootName;
            }
            case PbsAst.GroupExpr groupExpr -> hostOwnerHint(groupExpr.expression());
            default -> "";
        };
    }

    private String resolveExpressionOwner(
            final PbsAst.Expression expression,
            final PbsExecutableLoweringContext context) {
        if (expression == null) {
            return "";
        }
            return switch (expression) {
            case PbsAst.IdentifierExpr identifierExpr ->
                    resolveIdentifierOwner(identifierExpr.name(), context);
            case PbsAst.CallExpr callExpr -> resolveCallReturnOwner(callExpr, context);
            case PbsAst.MemberExpr memberExpr -> resolveExpressionOwner(memberExpr.receiver(), context);
            case PbsAst.GroupExpr groupExpr -> resolveExpressionOwner(groupExpr.expression(), context);
            default -> "";
        };
    }

    private String resolveIdentifierOwner(
            final String identifier,
            final PbsExecutableLoweringContext context) {
        final var localOwner = context.resolveLocalOwner(identifier);
        if (!localOwner.isBlank()) {
            return localOwner;
        }
        return context.builtinConstOwnerByNameId().getOrDefault(context.nameTable().register(identifier), "");
    }

    private String resolveCallReturnOwner(
            final PbsAst.CallExpr callExpr,
            final PbsExecutableLoweringContext context) {
        final var calleeIdentity = resolveCalleeIdentity(callExpr.callee(), context.nameTable());
        if (calleeIdentity == null) {
            return "";
        }
        final var receiverOwner = resolveReceiverOwnerForCallee(callExpr.callee(), context);
        if (!receiverOwner.isBlank()) {
            final var key = new PbsIntrinsicOwnerMethodKey(receiverOwner, calleeIdentity.memberSourceMethodName());
            final var intrinsicCandidates = context.intrinsicByOwnerAndMethod().getOrDefault(key, List.of());
            if (intrinsicCandidates.size() == 1) {
                final var intrinsic = intrinsicCandidates.getFirst();
                return context.intrinsicReturnOwnerByCanonical().getOrDefault(
                        new PbsIntrinsicCanonicalKey(intrinsic.canonicalName(), intrinsic.canonicalVersion()),
                        "");
            }
        }
        final var callableCandidates = resolveCallableCandidates(callExpr, calleeIdentity, context);
        if (callableCandidates.size() != 1) {
            return "";
        }
        return context.returnOwnerByCallableId().getOrDefault(callableCandidates.getFirst(), "");
    }

    private PbsCalleeIdentity resolveCalleeIdentity(
            final PbsAst.Expression callee,
            final NameTable nameTable) {
        return switch (callee) {
            case PbsAst.IdentifierExpr identifierExpr ->
                    new PbsCalleeIdentity(
                            nameTable.register(identifierExpr.name()),
                            null,
                            identifierExpr.name(),
                            nameTable.register(identifierExpr.name()));
            case PbsAst.MemberExpr memberExpr -> resolveMemberCalleeIdentity(memberExpr, nameTable);
            case PbsAst.BindExpr bindExpr ->
                    new PbsCalleeIdentity(
                            nameTable.register(bindExpr.functionName()),
                            null,
                            bindExpr.functionName(),
                            nameTable.register(bindExpr.functionName()));
            case PbsAst.GroupExpr groupExpr -> resolveCalleeIdentity(groupExpr.expression(), nameTable);
            default -> null;
        };
    }

    private PbsCalleeIdentity resolveMemberCalleeIdentity(
            final PbsAst.MemberExpr memberExpr,
            final NameTable nameTable) {
        final var memberName = memberExpr.memberName();
        final var rootName = memberRootName(memberExpr.receiver());
        final var qualified = rootName == null || rootName.isBlank()
                ? memberName
                : rootName + "." + memberName;
        final NameId secondaryNameId;
        if (qualified.equals(memberName)) {
            secondaryNameId = null;
        } else {
            secondaryNameId = nameTable.register(memberName);
        }
        return new PbsCalleeIdentity(
                nameTable.register(qualified),
                secondaryNameId,
                memberName,
                nameTable.register(memberName));
    }

    private String memberRootName(final PbsAst.Expression expression) {
        return switch (expression) {
            case PbsAst.IdentifierExpr identifierExpr -> identifierExpr.name();
            case PbsAst.MemberExpr memberExpr -> memberRootName(memberExpr.receiver());
            case PbsAst.CallExpr callExpr -> memberRootName(callExpr.callee());
            case PbsAst.GroupExpr groupExpr -> memberRootName(groupExpr.expression());
            default -> null;
        };
    }

    private int implicitReceiverArgSlots(
            final PbsAst.Expression callee,
            final PbsExecutableLoweringContext context) {
        if (!(callee instanceof PbsAst.MemberExpr memberExpr)) {
            return 0;
        }
        return receiverProducesRuntimeValue(memberExpr.receiver(), context) ? 1 : 0;
    }

    private boolean receiverProducesRuntimeValue(
            final PbsAst.Expression receiver,
            final PbsExecutableLoweringContext context) {
        return switch (receiver) {
            case PbsAst.CallExpr ignored -> true;
            case PbsAst.IdentifierExpr identifierExpr ->
                    context.resolveLocalSlot(identifierExpr.name()) != null
                            || context.resolveGlobalSlot(identifierExpr.name()) != null;
            case PbsAst.GroupExpr groupExpr -> receiverProducesRuntimeValue(groupExpr.expression(), context);
            default -> false;
        };
    }

    private void reportUnsupportedLowering(
            final String message,
            final p.studio.compiler.source.Span span,
            final PbsExecutableLoweringContext context) {
        context.diagnostics().error(
                DiagnosticPhase.STATIC_SEMANTICS,
                PbsSemanticsErrors.E_SEM_EXEC_LOWERING_UNRESOLVED_CALLEE.name(),
                PbsSemanticsErrors.E_SEM_EXEC_LOWERING_UNRESOLVED_CALLEE.name(),
                Map.of(),
                message,
                span);
    }
}
