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.Diagnostics;

final class PbsFlowStatementAnalyzer {
    @FunctionalInterface
    interface AssignmentStatementAnalyzer {
        void analyze(PbsAst.AssignStatement assignStatement, PbsFlowBodyContext context);
    }

    @FunctionalInterface
    interface InlineHintCollector {
        void collect(PbsAst.LetStatement letStatement, TypeView inferredType);

        static InlineHintCollector noop() {
            return (letStatement, inferredType) -> {
            };
        }
    }

    private final PbsFlowTypeOps typeOps;
    private final PbsFlowExpressionAnalyzer expressionAnalyzer;
    private final AssignmentStatementAnalyzer assignmentStatementAnalyzer;
    private final InlineHintCollector inlineHintCollector;

    PbsFlowStatementAnalyzer(
            final PbsFlowTypeOps typeOps,
            final PbsFlowExpressionAnalyzer expressionAnalyzer,
            final AssignmentStatementAnalyzer assignmentStatementAnalyzer,
            final InlineHintCollector inlineHintCollector) {
        this.typeOps = typeOps;
        this.expressionAnalyzer = expressionAnalyzer;
        this.assignmentStatementAnalyzer = assignmentStatementAnalyzer;
        this.inlineHintCollector = inlineHintCollector;
    }

    TypeView analyzeBlock(
            final PbsAst.Block block,
            final PbsFlowBodyContext context,
            final boolean valueContext) {
        final var scope = context.scope().copy();
        final var scopedContext = context.withScope(scope);
        for (final var statement : block.statements()) {
            analyzeStatement(statement, scopedContext);
        }
        if (block.tailExpression() == null) {
            return TypeView.unit();
        }
        return expressionAnalyzer.analyzeExpression(
                block.tailExpression(),
                scope,
                null,
                context.returnType(),
                context.resultErrorName(),
                context.receiverType(),
                context.model(),
                context.diagnostics(),
                ExprUse.VALUE,
                valueContext,
                this::analyzeNestedBlock).type();
    }

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

    private void analyzeStatement(
            final PbsAst.Statement statement,
            final PbsFlowBodyContext context) {
        final var scope = context.scope();
        final var returnType = context.returnType();
        final var resultErrorName = context.resultErrorName();
        final var receiverType = context.receiverType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        if (statement instanceof PbsAst.LetStatement letStatement) {
            final var expected = letStatement.explicitType() == null
                    ? null
                    : typeOps.typeFromTypeRef(letStatement.explicitType(), model, receiverType);
            final var initializer = expressionAnalyzer.analyzeExpression(
                    letStatement.initializer(),
                    scope,
                    expected,
                    returnType,
                    resultErrorName,
                    receiverType,
                    model,
                    diagnostics,
                    ExprUse.VALUE,
                    true,
                    this::analyzeNestedBlock);
            final var resolvedType = expected == null ? initializer.type() : expected;
            scope.bind(letStatement.name(), resolvedType, !letStatement.isConst());
            if (letStatement.explicitType() == null) {
                inlineHintCollector.collect(letStatement, resolvedType);
            }
            return;
        }
        if (statement instanceof PbsAst.ReturnStatement returnStatement) {
            if (returnStatement.value() != null) {
                analyzeReturnStatement(returnStatement.value(), context);
            }
            return;
        }
        if (statement instanceof PbsAst.IfStatement ifStatement) {
            final var condition = expressionAnalyzer.analyzeExpression(
                    ifStatement.condition(),
                    scope,
                    TypeView.bool(),
                    returnType,
                    resultErrorName,
                    receiverType,
                    model,
                    diagnostics,
                    ExprUse.VALUE,
                    true,
                    this::analyzeNestedBlock).type();
            if (!typeOps.isBool(condition)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_IF_NON_BOOL_CONDITION.name(),
                        "If statement condition must have bool type",
                        ifStatement.condition().span());
            }
            analyzeBlock(ifStatement.thenBlock(), context, false);
            if (ifStatement.elseBlock() != null) {
                analyzeBlock(ifStatement.elseBlock(), context, false);
            }
            if (ifStatement.elseIf() != null) {
                analyzeStatement(ifStatement.elseIf(), context);
            }
            return;
        }
        if (statement instanceof PbsAst.ForStatement(
                String iteratorName, PbsAst.TypeRef type, PbsAst.Expression fromExpression,
                PbsAst.Expression untilExpression, PbsAst.Expression stepExpression, PbsAst.Block body,
                Span span
        )) {
            analyzeForStatement(
                    iteratorName,
                    type,
                    fromExpression,
                    untilExpression,
                    stepExpression,
                    body,
                    span,
                    context);
            return;
        }
        if (statement instanceof PbsAst.WhileStatement whileStatement) {
            final var condition = expressionAnalyzer.analyzeExpression(
                    whileStatement.condition(),
                    scope,
                    TypeView.bool(),
                    returnType,
                    resultErrorName,
                    receiverType,
                    model,
                    diagnostics,
                    ExprUse.VALUE,
                    true,
                    this::analyzeNestedBlock).type();
            if (!typeOps.isBool(condition)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_WHILE_NON_BOOL_CONDITION.name(),
                        "While condition must have bool type",
                        whileStatement.condition().span());
            }
            analyzeBlock(whileStatement.body(), context, false);
            return;
        }
        if (statement instanceof PbsAst.ExpressionStatement expressionStatement) {
            final var expressionType = expressionAnalyzer.analyzeExpression(
                    expressionStatement.expression(),
                    scope,
                    null,
                    returnType,
                    resultErrorName,
                    receiverType,
                    model,
                    diagnostics,
                    ExprUse.VALUE,
                    false,
                    this::analyzeNestedBlock).type();
            if (typeOps.isMaterializedValue(expressionType)) {
                Diagnostics.warning(diagnostics,
                        PbsSemanticsErrors.W_SEM_IGNORED_VALUE.name(),
                        "Ignored materialized value in expression statement",
                        expressionStatement.span());
            }
            return;
        }
        if (statement instanceof PbsAst.AssignStatement assignStatement) {
            assignmentStatementAnalyzer.analyze(assignStatement, context);
        }
    }

    private void analyzeForStatement(
            final String iteratorName,
            final PbsAst.TypeRef type,
            final PbsAst.Expression fromExpression,
            final PbsAst.Expression untilExpression,
            final PbsAst.Expression stepExpression,
            final PbsAst.Block body,
            final Span span,
            final PbsFlowBodyContext context) {
        final var scope = context.scope();
        final var returnType = context.returnType();
        final var resultErrorName = context.resultErrorName();
        final var receiverType = context.receiverType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var iteratorType = typeOps.typeFromTypeRef(type, model, receiverType);
        final var fromType = expressionAnalyzer.analyzeExpression(
                fromExpression,
                scope,
                iteratorType,
                returnType,
                resultErrorName,
                receiverType,
                model,
                diagnostics,
                ExprUse.VALUE,
                true,
                this::analyzeNestedBlock).type();
        final var untilType = expressionAnalyzer.analyzeExpression(
                untilExpression,
                scope,
                iteratorType,
                returnType,
                resultErrorName,
                receiverType,
                model,
                diagnostics,
                ExprUse.VALUE,
                true,
                this::analyzeNestedBlock).type();
        if (!typeOps.compatible(fromType, iteratorType) || !typeOps.compatible(untilType, iteratorType)) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_FOR_TYPE_MISMATCH.name(),
                    "For-loop bounds must match declared iterator type",
                    span);
        }
        if (stepExpression != null) {
            final var stepType = expressionAnalyzer.analyzeExpression(
                    stepExpression,
                    scope,
                    iteratorType,
                    returnType,
                    resultErrorName,
                    receiverType,
                    model,
                    diagnostics,
                    ExprUse.VALUE,
                    true,
                    this::analyzeNestedBlock).type();
            if (!typeOps.compatible(stepType, iteratorType)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_FOR_TYPE_MISMATCH.name(),
                        "For-loop step must match declared iterator type",
                        stepExpression.span());
            }
        }
        final var bodyScope = scope.copy();
        bodyScope.bind(iteratorName, iteratorType);
        analyzeBlock(body, context.withScope(bodyScope), false);
    }

    private void analyzeReturnStatement(
            final PbsAst.Expression value,
            final PbsFlowBodyContext context) {
        final var root = unwrapGroup(value);
        if (root instanceof PbsAst.OkExpr okExpr) {
            analyzeReturnOk(okExpr, context);
            return;
        }
        if (root instanceof PbsAst.ErrExpr errExpr) {
            analyzeReturnErr(errExpr, context);
            return;
        }
        expressionAnalyzer.analyzeExpression(
                value,
                context.scope(),
                context.returnType(),
                context.returnType(),
                context.resultErrorName(),
                context.receiverType(),
                context.model(),
                context.diagnostics(),
                ExprUse.VALUE,
                true,
                this::analyzeNestedBlock);
    }

    private void analyzeReturnOk(
            final PbsAst.OkExpr okExpr,
            final PbsFlowBodyContext context) {
        final var returnType = context.returnType();
        final var resultErrorName = context.resultErrorName();
        final var diagnostics = context.diagnostics();
        if (returnType.kind() != Kind.RESULT || resultErrorName == null) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_FLOW_INVALID_POSITION.name(),
                    "'ok(...)' is only allowed when returning from a result callable",
                    okExpr.span());
            expressionAnalyzer.analyzeExpression(
                    okExpr.value(),
                    context.scope(),
                    null,
                    returnType,
                    resultErrorName,
                    context.receiverType(),
                    context.model(),
                    diagnostics,
                    ExprUse.VALUE,
                    true,
                    this::analyzeNestedBlock);
            return;
        }

        final var payloadType = returnType.inner();
        final var actualType = expressionAnalyzer.analyzeExpression(
                okExpr.value(),
                context.scope(),
                payloadType,
                returnType,
                resultErrorName,
                context.receiverType(),
                context.model(),
                diagnostics,
                ExprUse.VALUE,
                true,
                this::analyzeNestedBlock).type();
        if (!typeOps.compatible(actualType, payloadType)) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_OK_PAYLOAD_MISMATCH.name(),
                    "Payload in 'ok(...)' is incompatible with result payload type",
                    okExpr.value().span());
        }
    }

    private void analyzeReturnErr(
            final PbsAst.ErrExpr errExpr,
            final PbsFlowBodyContext context) {
        final var returnType = context.returnType();
        final var resultErrorName = context.resultErrorName();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        if (returnType.kind() != Kind.RESULT || resultErrorName == null) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_FLOW_INVALID_POSITION.name(),
                    "'err(...)' is only allowed when returning from a result callable",
                    errExpr.span());
            return;
        }

        if (!matchesTargetError(errExpr.errorPath(), resultErrorName, model)) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_RESULT_ERROR_LABEL_INVALID.name(),
                    "Error label in 'err(...)' does not match enclosing result error type",
                    errExpr.errorPath().span());
        }
    }

    private PbsAst.Expression unwrapGroup(final PbsAst.Expression expression) {
        if (expression instanceof PbsAst.GroupExpr groupExpr) {
            return unwrapGroup(groupExpr.expression());
        }
        return expression;
    }

    private boolean matchesTargetError(
            final PbsAst.ErrorPath path,
            final String resultErrorName,
            final Model model) {
        if (path == null || resultErrorName == null || path.segments().size() != 2) {
            return false;
        }
        final var errorName = path.segments().getFirst();
        final var caseName = path.segments().get(1);
        final var targetCases = model.errors.get(errorName);
        return resultErrorName.equals(errorName) && targetCases != null && targetCases.contains(caseName);
    }
}
