package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.ExprResult;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Kind;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TupleField;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TypeView;
import p.studio.compiler.source.diagnostics.Diagnostics;

import java.util.ArrayList;

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

    @FunctionalInterface
    interface BlockAnalysisDelegate {
        TypeView analyze(PbsAst.Block block, PbsFlowExpressionContext context, boolean valueContext);
    }

    private final PbsFlowTypeOps typeOps;
    private final ExpressionAnalyzerDelegate expressionAnalyzerDelegate;
    private final BlockAnalysisDelegate blockAnalysisDelegate;

    PbsFlowStructuralExpressionAnalyzer(
            final PbsFlowTypeOps typeOps,
            final ExpressionAnalyzerDelegate expressionAnalyzerDelegate,
            final BlockAnalysisDelegate blockAnalysisDelegate) {
        this.typeOps = typeOps;
        this.expressionAnalyzerDelegate = expressionAnalyzerDelegate;
        this.blockAnalysisDelegate = blockAnalysisDelegate;
    }

    ExprResult analyze(
            final PbsAst.Expression expression,
            final PbsFlowExpressionContext context) {
        final var expectedType = context.expectedType();
        final var receiverType = context.receiverType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();

        if (expression instanceof PbsAst.IntLiteralExpr) {
            return ExprResult.type(TypeView.intType());
        }
        if (expression instanceof PbsAst.FloatLiteralExpr) {
            return ExprResult.type(TypeView.floatType());
        }
        if (expression instanceof PbsAst.BoundedLiteralExpr) {
            return ExprResult.type(TypeView.intType());
        }
        if (expression instanceof PbsAst.StringLiteralExpr) {
            return ExprResult.type(TypeView.str());
        }
        if (expression instanceof PbsAst.BoolLiteralExpr) {
            return ExprResult.type(TypeView.bool());
        }
        if (expression instanceof PbsAst.UnitExpr) {
            return ExprResult.type(TypeView.unit());
        }
        if (expression instanceof PbsAst.ThisExpr thisExpr) {
            if (receiverType == null) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_THIS_CONTEXT.name(),
                        "Invalid 'this' usage outside struct/service methods and constructors",
                        thisExpr.span());
                return ExprResult.type(TypeView.unknown());
            }
            return ExprResult.type(receiverType);
        }
        if (expression instanceof PbsAst.IdentifierExpr identifierExpr) {
            if ("assets".equals(identifierExpr.name()) && model.hasAssetNamespaceRoot()) {
                return ExprResult.type(TypeView.assetNamespace("assets"));
            }

            final var localType = context.scope().resolve(identifierExpr.name());
            if (localType != null) {
                return ExprResult.type(localType);
            }

            final var serviceType = model.serviceSingletons.get(identifierExpr.name());
            if (serviceType != null) {
                return ExprResult.type(serviceType);
            }

            final var constType = model.constTypes.get(identifierExpr.name());
            if (constType != null) {
                return ExprResult.type(constType);
            }

            final var globalType = model.globalTypes.get(identifierExpr.name());
            if (globalType != null) {
                return ExprResult.type(globalType);
            }

            final var callbackSignature = model.callbacks.get(identifierExpr.name());
            if (callbackSignature != null) {
                return ExprResult.type(TypeView.callback(identifierExpr.name(), callbackSignature.inputTypes(), callbackSignature.outputType()));
            }

            final var callables = model.topLevelCallables.get(identifierExpr.name());
            if (callables != null && !callables.isEmpty()) {
                return ExprResult.callables(callables, false);
            }

            if (model.enums.containsKey(identifierExpr.name())) {
                return ExprResult.type(TypeView.typeRef(identifierExpr.name()));
            }

            return ExprResult.type(TypeView.unknown());
        }
        if (expression instanceof PbsAst.GroupExpr groupExpr) {
            return expressionAnalyzerDelegate.analyze(groupExpr.expression(), context);
        }
        if (expression instanceof PbsAst.TupleExpr tupleExpr) {
            final var fields = new ArrayList<TupleField>(tupleExpr.items().size());
            for (final var item : tupleExpr.items()) {
                final var value = analyzeValueExpression(item.expression(), context, null).type();
                fields.add(new TupleField(item.label(), value));
            }
            return ExprResult.type(TypeView.tuple(fields));
        }
        if (expression instanceof PbsAst.BlockExpr blockExpr) {
            return ExprResult.type(blockAnalysisDelegate.analyze(blockExpr.block(), context, context.valueContext()));
        }
        if (expression instanceof PbsAst.UnaryExpr unaryExpr) {
            final var operand = analyzeValueExpression(unaryExpr.expression(), context, null).type();
            if ("!".equals(unaryExpr.operator()) || "not".equals(unaryExpr.operator())) {
                return ExprResult.type(TypeView.bool());
            }
            if ("-".equals(unaryExpr.operator()) && (typeOps.isInt(operand) || typeOps.isFloat(operand))) {
                return ExprResult.type(operand);
            }
            return ExprResult.type(TypeView.unknown());
        }
        if (expression instanceof PbsAst.BinaryExpr binaryExpr) {
            final var left = analyzeValueExpression(binaryExpr.left(), context, null).type();
            final var right = analyzeValueExpression(binaryExpr.right(), context, null).type();
            return ExprResult.type(typeOps.inferBinaryResult(binaryExpr.operator(), left, right));
        }
        if (expression instanceof PbsAst.NoneExpr noneExpr) {
            if (expectedType == null || expectedType.kind() != Kind.OPTIONAL) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_NONE_WITHOUT_EXPECTED_OPTIONAL.name(),
                        "'none' requires an expected optional type context",
                        noneExpr.span());
                return ExprResult.type(TypeView.unknown());
            }
            return ExprResult.type(expectedType);
        }
        if (expression instanceof PbsAst.SomeExpr someExpr) {
            final var payloadExpected = expectedType != null && expectedType.kind() == Kind.OPTIONAL
                    ? expectedType.inner()
                    : null;
            final var payload = analyzeValueExpression(someExpr.value(), context, payloadExpected).type();
            return ExprResult.type(TypeView.optional(payloadExpected == null ? payload : payloadExpected));
        }
        if (expression instanceof PbsAst.NewExpr newExpr) {
            if (model.structs.containsKey(newExpr.typeName())) {
                return ExprResult.type(TypeView.struct(newExpr.typeName()));
            }
            return ExprResult.type(TypeView.unknown());
        }
        if (expression instanceof PbsAst.AsExpr asExpr) {
            analyzeValueExpression(asExpr.expression(), context, null);
            if (model.contracts.containsKey(asExpr.contractName())) {
                return ExprResult.type(TypeView.contract(asExpr.contractName()));
            }
            return ExprResult.type(TypeView.unknown());
        }
        if (expression instanceof PbsAst.OkExpr okExpr) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_FLOW_INVALID_POSITION.name(),
                    "'ok(...)' is only allowed in result return flow and handle arm terminals",
                    okExpr.span());
            analyzeValueExpression(okExpr.value(), context, null);
            return ExprResult.type(TypeView.unknown());
        }
        if (expression instanceof PbsAst.ErrExpr errExpr) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_FLOW_INVALID_POSITION.name(),
                    "'err(...)' is only allowed in result return flow and handle arm terminals",
                    errExpr.span());
            return ExprResult.type(TypeView.unknown());
        }

        return null;
    }

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