package p.studio.compiler.pbs.semantics;

import p.studio.compiler.messages.FESurfaceContext;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Model;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Scope;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TypeView;
import p.studio.compiler.services.PbsInlineHintSurface;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;

final class PbsFlowBodyAnalyzer {
    private final PbsFlowTypeOps typeOps = new PbsFlowTypeOps();
    private final PbsFlowExpressionAnalyzer expressionAnalyzer = new PbsFlowExpressionAnalyzer(typeOps);
    private final PbsFlowCompletionAnalyzer completionAnalyzer = new PbsFlowCompletionAnalyzer();
    private final PbsFlowAssignmentAnalyzer assignmentAnalyzer = new PbsFlowAssignmentAnalyzer(
            typeOps,
            expressionAnalyzer,
            this::analyzeBlock);
    private final PbsFlowCallableBodyAnalyzer callableBodyAnalyzer = new PbsFlowCallableBodyAnalyzer(
            typeOps,
            completionAnalyzer,
            this::analyzeBlock);
    private PbsFlowStatementAnalyzer.InlineHintCollector inlineHintCollector = PbsFlowStatementAnalyzer.InlineHintCollector.noop();

    public void validate(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final FESurfaceContext feSurfaceContext,
            final DiagnosticSink diagnostics) {
        final var previousCollector = inlineHintCollector;
        inlineHintCollector = PbsFlowStatementAnalyzer.InlineHintCollector.noop();
        try {
            final var model = Model.from(ast, supplementalTopDecls, feSurfaceContext, diagnostics);

            for (final var topDecl : ast.topDecls()) {
                validateTopDecl(topDecl, model, diagnostics);
            }
        } finally {
            inlineHintCollector = previousCollector;
        }
    }

    public ReadOnlyList<PbsInlineHintSurface> collectInlineHints(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final FESurfaceContext feSurfaceContext,
            final DiagnosticSink diagnostics) {
        final var collector = new InlineHintsCollector();
        final var previousCollector = inlineHintCollector;
        inlineHintCollector = collector;
        try {
            final var model = Model.from(ast, supplementalTopDecls, feSurfaceContext, diagnostics);

            for (final var topDecl : ast.topDecls()) {
                validateTopDecl(topDecl, model, diagnostics);
            }
        } finally {
            inlineHintCollector = previousCollector;
        }
        return collector.hints();
    }

    private PbsFlowStatementAnalyzer statementAnalyzer() {
        return new PbsFlowStatementAnalyzer(
                typeOps,
                expressionAnalyzer,
                assignmentAnalyzer::analyzeAssignmentStatement,
                inlineHintCollector);
    }

    private static final class InlineHintsCollector implements PbsFlowStatementAnalyzer.InlineHintCollector {
        private static final String TYPE_HINT_CATEGORY = "type";

        private final ArrayList<PbsInlineHintSurface> hints = new ArrayList<>();
        private final PbsInlineHintSurfaceFormatter formatter = new PbsInlineHintSurfaceFormatter();

        @Override
        public void collect(final PbsAst.LetStatement letStatement, final TypeView inferredType) {
            final var label = formatter.format(inferredType);
            if (label == null || letStatement.nameSpan() == null || letStatement.nameSpan().isNone()) {
                return;
            }
            hints.add(new PbsInlineHintSurface(letStatement.nameSpan(), label, TYPE_HINT_CATEGORY));
        }

        ReadOnlyList<PbsInlineHintSurface> hints() {
            return ReadOnlyList.wrap(hints);
        }
    }

    private void validateTopDecl(
            final PbsAst.TopDecl topDecl,
            final Model model,
            final DiagnosticSink diagnostics) {
        if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
            validateFunction(functionDecl, model, diagnostics);
            return;
        }
        if (topDecl instanceof PbsAst.StructDecl structDecl) {
            validateStruct(structDecl, model, diagnostics);
            return;
        }
        if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
            validateService(serviceDecl, model, diagnostics);
        }
    }

    private void validateFunction(
            final PbsAst.FunctionDecl functionDecl,
            final Model model,
            final DiagnosticSink diagnostics) {
        validateCallableBody(
                functionDecl.parameters(),
                functionDecl.body(),
                typeOps.callableReturnType(
                        functionDecl.returnKind(),
                        functionDecl.returnType(),
                        functionDecl.resultErrorType(),
                        model),
                errorName(functionDecl.resultErrorType()),
                null,
                model,
                diagnostics);
    }

    private void validateStruct(
            final PbsAst.StructDecl structDecl,
            final Model model,
            final DiagnosticSink diagnostics) {
        final var receiverType = TypeView.struct(structDecl.name());
        for (final var method : structDecl.methods()) {
            validateCallableBody(
                    method.parameters(),
                    method.body(),
                    typeOps.callableReturnType(
                            method.returnKind(),
                            method.returnType(),
                            method.resultErrorType(),
                            model),
                    errorName(method.resultErrorType()),
                    receiverType,
                    model,
                    diagnostics);
        }
        for (final var ctor : structDecl.ctors()) {
            validateCallableBody(
                    ctor.parameters(),
                    ctor.body(),
                    TypeView.unit(),
                    null,
                    receiverType,
                    model,
                    diagnostics);
        }
    }

    private void validateService(
            final PbsAst.ServiceDecl serviceDecl,
            final Model model,
            final DiagnosticSink diagnostics) {
        final var receiverType = TypeView.service(serviceDecl.name());
        for (final var method : serviceDecl.methods()) {
            validateCallableBody(
                    method.parameters(),
                    method.body(),
                    typeOps.callableReturnType(
                            method.returnKind(),
                            method.returnType(),
                            method.resultErrorType(),
                            model),
                    errorName(method.resultErrorType()),
                    receiverType,
                    model,
                    diagnostics);
        }
    }

    private PbsFlowBodyContext newCallableContext(
            final TypeView returnType,
            final String resultErrorName,
            final TypeView receiverType,
            final Model model,
            final DiagnosticSink diagnostics) {
        return new PbsFlowBodyContext(
                new Scope(),
                returnType,
                resultErrorName,
                receiverType,
                model,
                diagnostics);
    }

    private void validateCallableBody(
            final ReadOnlyList<PbsAst.Parameter> parameters,
            final PbsAst.Block body,
            final TypeView returnType,
            final String resultErrorName,
            final TypeView receiverType,
            final Model model,
            final DiagnosticSink diagnostics) {
        validateCallableBody(
                parameters,
                body,
                newCallableContext(
                        returnType,
                        resultErrorName,
                        receiverType,
                        model,
                        diagnostics));
    }

    private String errorName(final PbsAst.TypeRef errorType) {
        return errorType == null ? null : errorType.name();
    }

    private void validateCallableBody(
            final ReadOnlyList<PbsAst.Parameter> parameters,
            final PbsAst.Block body,
            final PbsFlowBodyContext context) {
        callableBodyAnalyzer.validateCallableBody(parameters, body, context);
    }

    private TypeView analyzeBlock(
            final PbsAst.Block block,
            final Scope outerScope,
            final TypeView returnType,
            final String resultErrorName,
            final TypeView receiverType,
            final Model model,
            final DiagnosticSink diagnostics,
            final boolean valueContext) {
        return analyzeBlock(
                block,
                new PbsFlowBodyContext(
                        outerScope,
                        returnType,
                        resultErrorName,
                        receiverType,
                        model,
                        diagnostics),
                valueContext);
    }

    private TypeView analyzeBlock(
            final PbsAst.Block block,
            final PbsFlowBodyContext context,
            final boolean valueContext) {
        return statementAnalyzer().analyzeBlock(block, context, valueContext);
    }
}
