package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Kind;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Model;
import p.studio.compiler.source.diagnostics.Diagnostics;

import java.util.HashSet;

final class PbsFlowCompletionAnalyzer {
    void validateCallableCompletion(
            final PbsAst.Block body,
            final PbsFlowBodyContext context) {
        final var returnType = context.returnType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var returnKind = returnType.kind();
        final var enforceResultCompletion = returnKind == Kind.RESULT;
        final var enforcePlainNonUnitCompletion = returnKind != Kind.RESULT
                && returnKind != Kind.UNIT
                && returnKind != Kind.OPTIONAL
                && returnKind != Kind.UNKNOWN;

        if (!enforceResultCompletion && !enforcePlainNonUnitCompletion) {
            return;
        }
        if (blockAlwaysReturns(body, model)) {
            return;
        }

        if (enforceResultCompletion) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name(),
                    "Possible fallthrough in result callable body; all paths must return explicitly",
                    body.span());
            return;
        }

        Diagnostics.error(diagnostics,
                PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_NON_UNIT.name(),
                "Possible fallthrough in plain non-unit callable body; all paths must return explicitly",
                body.span());
    }

    private boolean blockAlwaysReturns(
            final PbsAst.Block block,
            final Model model) {
        if (block == null) {
            return false;
        }

        for (final var statement : block.statements()) {
            if (statementAlwaysReturns(statement, model)) {
                return true;
            }
        }
        return expressionAlwaysReturns(block.tailExpression(), model);
    }

    private boolean statementAlwaysReturns(
            final PbsAst.Statement statement,
            final Model model) {
        if (statement instanceof PbsAst.ReturnStatement) {
            return true;
        }
        if (statement instanceof PbsAst.IfStatement ifStatement) {
            final var thenReturns = blockAlwaysReturns(ifStatement.thenBlock(), model);
            final boolean elseReturns;
            if (ifStatement.elseBlock() != null) {
                elseReturns = blockAlwaysReturns(ifStatement.elseBlock(), model);
            } else {
                elseReturns = ifStatement.elseIf() != null && statementAlwaysReturns(ifStatement.elseIf(), model);
            }
            return thenReturns && elseReturns;
        }
        if (statement instanceof PbsAst.ExpressionStatement expressionStatement) {
            return expressionAlwaysReturns(expressionStatement.expression(), model);
        }
        return false;
    }

    private boolean expressionAlwaysReturns(
            final PbsAst.Expression expression,
            final Model model) {
        return switch (expression) {
            case PbsAst.GroupExpr groupExpr -> expressionAlwaysReturns(groupExpr.expression(), model);
            case PbsAst.BlockExpr blockExpr -> blockAlwaysReturns(blockExpr.block(), model);
            case PbsAst.IfExpr ifExpr -> blockAlwaysReturns(ifExpr.thenBlock(), model)
                    && expressionAlwaysReturns(ifExpr.elseExpression(), model);
            case PbsAst.SwitchExpr switchExpr -> switchAlwaysReturns(switchExpr, model);
            case PbsAst.HandleExpr handleExpr -> handleAlwaysReturns(handleExpr, model);
            case null, default -> false;
        };
    }

    private boolean handleAlwaysReturns(
            final PbsAst.HandleExpr handleExpr,
            final Model model) {
        if (expressionAlwaysReturns(handleExpr.value(), model)) {
            return true;
        }
        if (handleExpr.arms().isEmpty()) {
            return false;
        }
        for (final var arm : handleExpr.arms()) {
            if (!handleArmAlwaysReturns(arm, model)) {
                return false;
            }
        }
        return true;
    }

    private boolean handleArmAlwaysReturns(
            final PbsAst.HandleArm arm,
            final Model model) {
        if (arm.remapTarget() != null) {
            return true;
        }
        if (arm.block() == null) {
            return false;
        }
        final var terminal = unwrapGroup(arm.block().tailExpression());
        if (terminal instanceof PbsAst.ErrExpr errExpr) {
            return isKnownErrorPath(errExpr.errorPath(), model);
        }
        return false;
    }

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

    private boolean switchAlwaysReturns(
            final PbsAst.SwitchExpr switchExpr,
            final Model model) {
        if (switchExpr.arms().isEmpty()) {
            return false;
        }
        if (!switchIsExhaustive(switchExpr, model)) {
            return false;
        }
        for (final var arm : switchExpr.arms()) {
            if (!blockAlwaysReturns(arm.block(), model)) {
                return false;
            }
        }
        return true;
    }

    private boolean switchIsExhaustive(
            final PbsAst.SwitchExpr switchExpr,
            final Model model) {
        for (final var arm : switchExpr.arms()) {
            if (arm.pattern() instanceof PbsAst.WildcardSwitchPattern) {
                return true;
            }
        }

        String enumName = null;
        final var coveredCases = new HashSet<String>();
        for (final var arm : switchExpr.arms()) {
            if (!(arm.pattern() instanceof PbsAst.EnumCaseSwitchPattern enumCaseSwitchPattern)) {
                return false;
            }
            final var segments = enumCaseSwitchPattern.path().segments();
            if (segments.size() != 2) {
                return false;
            }
            if (enumName == null) {
                enumName = segments.getFirst();
            } else if (!enumName.equals(segments.getFirst())) {
                return false;
            }
            coveredCases.add(segments.get(1));
        }
        if (enumName == null) {
            return false;
        }

        final var enumCases = model.enums.get(enumName);
        return enumCases != null && coveredCases.containsAll(enumCases);
    }

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