package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.*;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.diagnostics.Diagnostics;

import java.util.ArrayList;

final class PbsFlowCallableResolutionAnalyzer {
    @FunctionalInterface
    interface ExpressionAnalyzerDelegate {
        ExprResult analyze(PbsAst.Expression expression, PbsFlowExpressionContext context);
    }

    private final PbsFlowTypeOps typeOps;
    private final ExpressionAnalyzerDelegate expressionAnalyzerDelegate;

    PbsFlowCallableResolutionAnalyzer(
            final PbsFlowTypeOps typeOps,
            final ExpressionAnalyzerDelegate expressionAnalyzerDelegate) {
        this.typeOps = typeOps;
        this.expressionAnalyzerDelegate = expressionAnalyzerDelegate;
    }

    ExprResult analyzeCallExpression(
            final PbsAst.CallExpr callExpr,
            final PbsFlowExpressionContext context) {
        final var callee = analyzeTargetExpression(callExpr.callee(), context, ExprUse.CALL_TARGET);

        final TypeView argumentType;
        if (callExpr.arguments().isEmpty()) {
            argumentType = TypeView.unit();
        } else if (callExpr.arguments().size() == 1) {
            argumentType = analyzeValueExpression(callExpr.arguments().getFirst(), context, null).type();
        } else {
            final var fields = new ArrayList<TupleField>(callExpr.arguments().size());
            for (final var argument : callExpr.arguments()) {
                final var itemType = analyzeValueExpression(argument, context, null).type();
                fields.add(new TupleField(null, itemType));
            }
            argumentType = TypeView.tuple(fields);
        }

        return resolveCallableApplication(
                callExpr.span(),
                callExpr.callee().span(),
                callee,
                argumentType,
                context.expectedType(),
                context.diagnostics());
    }

    ExprResult analyzeApplyExpression(
            final PbsAst.ApplyExpr applyExpr,
            final PbsFlowExpressionContext context) {
        final var callee = analyzeTargetExpression(applyExpr.callee(), context, ExprUse.APPLY_TARGET);
        final var argument = analyzeValueExpression(applyExpr.argument(), context, null).type();

        return resolveCallableApplication(
                applyExpr.span(),
                applyExpr.callee().span(),
                callee,
                argument,
                context.expectedType(),
                context.diagnostics());
    }

    ExprResult analyzeMemberExpression(
            final PbsAst.MemberExpr memberExpr,
            final PbsFlowExpressionContext context) {
        final var assetReference = flattenAssetReference(memberExpr);
        if (assetReference != null) {
            return resolveAssetReference(assetReference, memberExpr, context);
        }

        final var receiver = analyzeValueExpression(memberExpr.receiver(), context, null).type();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var receiverType = context.receiverType();
        final var use = context.use();

        if (receiver.kind() == Kind.TYPE_REF) {
            final var enumCases = model.enums.get(receiver.name());
            if (enumCases != null && enumCases.contains(memberExpr.memberName())) {
                return ExprResult.type(TypeView.enumType(receiver.name()));
            }
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                    "Invalid type member access",
                    memberExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        if (receiver.kind() == Kind.TUPLE) {
            if (receiver.tupleFields().size() <= 1) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                        "Member projection is not defined for single-slot carrier values",
                        memberExpr.span());
                return ExprResult.type(TypeView.unknown());
            }
            for (final var field : receiver.tupleFields()) {
                if (memberExpr.memberName().equals(field.label())) {
                    return ExprResult.type(field.type());
                }
            }
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                    "Tuple label '%s' does not exist".formatted(memberExpr.memberName()),
                    memberExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        if (receiver.kind() == Kind.STRUCT) {
            final var struct = model.structs.get(receiver.name());
            if (struct == null) {
                return ExprResult.type(TypeView.unknown());
            }
            final var field = struct.fields().get(memberExpr.memberName());
            if (field != null) {
                if (!canReadStructField(field, receiver, memberExpr.receiver(), receiverType)) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_FIELD_READ_ACCESS_DENIED.name(),
                            "Read access to struct field '%s' is not permitted".formatted(memberExpr.memberName()),
                            memberExpr.span());
                    return ExprResult.type(TypeView.unknown());
                }
                return ExprResult.type(field.type());
            }
            final var methods = struct.methods().get(memberExpr.memberName());
            if (methods != null && !methods.isEmpty()) {
                if (use == ExprUse.VALUE) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_BARE_METHOD_EXTRACTION.name(),
                            "Bare method extraction is not allowed in PBS core",
                            memberExpr.span());
                    return ExprResult.type(TypeView.unknown());
                }
                return ExprResult.callables(methods, true);
            }
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                    "Struct member '%s' does not exist".formatted(memberExpr.memberName()),
                    memberExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        if (receiver.kind() == Kind.SERVICE) {
            final var service = model.services.get(receiver.name());
            if (service == null) {
                return ExprResult.type(TypeView.unknown());
            }
            final var methods = service.methods().get(memberExpr.memberName());
            if (methods != null && !methods.isEmpty()) {
                if (use == ExprUse.VALUE) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_BARE_METHOD_EXTRACTION.name(),
                            "Bare method extraction is not allowed in PBS core",
                            memberExpr.span());
                    return ExprResult.type(TypeView.unknown());
                }
                return ExprResult.callables(methods, true);
            }
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                    "Service member '%s' does not exist".formatted(memberExpr.memberName()),
                    memberExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        if (receiver.kind() == Kind.CONTRACT) {
            final var contract = model.contracts.get(receiver.name());
            if (contract == null) {
                return ExprResult.type(TypeView.unknown());
            }
            final var methods = contract.methods().get(memberExpr.memberName());
            if (methods != null && !methods.isEmpty()) {
                if (use == ExprUse.VALUE) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_BARE_METHOD_EXTRACTION.name(),
                            "Bare method extraction is not allowed in PBS core",
                            memberExpr.span());
                    return ExprResult.type(TypeView.unknown());
                }
                return ExprResult.callables(methods, true);
            }
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                    "Contract member '%s' does not exist".formatted(memberExpr.memberName()),
                    memberExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        Diagnostics.error(diagnostics,
                PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                "Invalid member access target",
                memberExpr.span());
        return ExprResult.type(TypeView.unknown());
    }

    ExprResult analyzeBindExpression(
            final PbsAst.BindExpr bindExpr,
            final PbsFlowExpressionContext context) {
        final var contextType = analyzeValueExpression(bindExpr.contextExpression(), context, null).type();
        final var expectedType = context.expectedType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        if (expectedType == null || expectedType.kind() != Kind.CALLBACK) {
            return ExprResult.type(TypeView.unknown());
        }

        final var topLevel = model.topLevelCallables.get(bindExpr.functionName());
        if (topLevel == null || topLevel.isEmpty()) {
            return ExprResult.type(expectedType);
        }
        final var compatible = new ArrayList<CallableSymbol>();
        for (final var candidate : topLevel) {
            if (candidate.inputTypes().size() != expectedType.callbackInputs().size() + 1) {
                continue;
            }
            if (!typeOps.compatible(contextType, candidate.inputTypes().getFirst())) {
                continue;
            }
            var ok = true;
            for (int i = 0; i < expectedType.callbackInputs().size(); i++) {
                if (!typeOps.compatible(expectedType.callbackInputs().get(i), candidate.inputTypes().get(i + 1))) {
                    ok = false;
                    break;
                }
            }
            if (!ok) {
                continue;
            }
            if (!typeOps.compatible(candidate.outputType(), expectedType.callbackOutput())) {
                continue;
            }
            compatible.add(candidate);
        }
        if (compatible.size() > 1) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_APPLY_AMBIGUOUS_OVERLOAD.name(),
                    "Bind target resolves ambiguously for callback type",
                    bindExpr.span());
        } else if (compatible.isEmpty()) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_APPLY_UNRESOLVED_OVERLOAD.name(),
                    "Bind target does not resolve for callback type",
                    bindExpr.span());
        }
        return ExprResult.type(expectedType);
    }

    private ExprResult resolveCallableApplication(
            final Span wholeSpan,
            final Span calleeSpan,
            final ExprResult callee,
            final TypeView argumentType,
            final TypeView expectedType,
            final DiagnosticSink diagnostics) {
        final var candidates = callee.callables();
        if (candidates.isEmpty()) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_APPLY_NON_CALLABLE_TARGET.name(),
                    "Apply/call target is not callable",
                    calleeSpan);
            return ExprResult.type(TypeView.unknown());
        }

        final var compatible = new ArrayList<CallableSymbol>();
        for (final var candidate : candidates) {
            if (typeOps.inputCompatible(candidate.inputTypes(), argumentType)) {
                compatible.add(candidate);
            }
        }

        if (compatible.isEmpty()) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_APPLY_UNRESOLVED_OVERLOAD.name(),
                    "No callable overload matches the provided argument",
                    wholeSpan);
            return ExprResult.type(TypeView.unknown());
        }

        if (compatible.size() > 1 && expectedType != null) {
            final var narrowed = compatible.stream()
                    .filter(candidate -> typeOps.compatible(candidate.outputType(), expectedType))
                    .toList();
            if (narrowed.size() == 1) {
                return ExprResult.type(narrowed.getFirst().outputType());
            }
            if (!narrowed.isEmpty()) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_APPLY_AMBIGUOUS_OVERLOAD.name(),
                        "Callable overload resolution remains ambiguous",
                        wholeSpan);
                return ExprResult.type(TypeView.unknown());
            }
        }

        if (compatible.size() > 1) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_APPLY_AMBIGUOUS_OVERLOAD.name(),
                    "Callable overload resolution is ambiguous",
                    wholeSpan);
            return ExprResult.type(TypeView.unknown());
        }

        return ExprResult.type(compatible.getFirst().outputType());
    }

    private boolean canReadStructField(
            final PbsFlowSemanticSupport.StructFieldInfo field,
            final TypeView receiver,
            final PbsAst.Expression receiverExpr,
            final TypeView currentReceiverType) {
        if (field.isPublic()) {
            return true;
        }
        return receiverExpr instanceof PbsAst.ThisExpr
                && currentReceiverType != null
                && currentReceiverType.kind() == Kind.STRUCT
                && receiver.kind() == Kind.STRUCT
                && receiver.name().equals(currentReceiverType.name());
    }

    private ExprResult resolveAssetReference(
            final String assetReference,
            final PbsAst.MemberExpr memberExpr,
            final PbsFlowExpressionContext context) {
        final var diagnostics = context.diagnostics();
        final var model = context.model();
        final var resolvedAssetId = model.resolveAssetId(assetReference);
        if (resolvedAssetId != null) {
            return ExprResult.type(TypeView.addressable(assetReference));
        }
        if (model.isAssetNamespace(assetReference)) {
            if (context.use() == ExprUse.VALUE) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                        "Asset namespace '%s' is not a terminal Addressable value".formatted(assetReference),
                        memberExpr.span());
                return ExprResult.type(TypeView.unknown());
            }
            return ExprResult.type(TypeView.assetNamespace(assetReference));
        }
        Diagnostics.error(diagnostics,
                PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                "Asset reference '%s' does not resolve in the backend-provided surface".formatted(assetReference),
                memberExpr.span());
        return ExprResult.type(TypeView.unknown());
    }

    private String flattenAssetReference(final PbsAst.Expression expression) {
        if (expression instanceof PbsAst.IdentifierExpr identifierExpr) {
            return "assets".equals(identifierExpr.name()) ? "assets" : null;
        }
        if (!(expression instanceof PbsAst.MemberExpr memberExpr)) {
            return null;
        }
        final var receiverPath = flattenAssetReference(memberExpr.receiver());
        if (receiverPath == null || receiverPath.isBlank()) {
            return null;
        }
        return receiverPath + "." + memberExpr.memberName();
    }

    private ExprResult analyzeValueExpression(
            final PbsAst.Expression expression,
            final PbsFlowExpressionContext context,
            final TypeView expectedType) {
        return expressionAnalyzerDelegate.analyze(
                expression,
                context.withExpectedType(expectedType)
                        .withUse(ExprUse.VALUE)
                        .withValueContext(true));
    }

    private ExprResult analyzeTargetExpression(
            final PbsAst.Expression expression,
            final PbsFlowExpressionContext context,
            final ExprUse use) {
        return expressionAnalyzerDelegate.analyze(
                expression,
                context.withExpectedType(null)
                        .withUse(use)
                        .withValueContext(true));
    }
}
