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.TypeView;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.diagnostics.Diagnostics;

import java.util.HashSet;
import java.util.Set;

final class PbsFlowControlExpressionAnalyzer {
    @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;

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

    TypeView analyzeSwitchExpression(
            final PbsAst.SwitchExpr switchExpr,
            final PbsFlowExpressionContext context) {
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var selector = analyzeValueExpression(switchExpr.selector(), context, null).type();

        final var selectorComparable = typeOps.isScalarComparable(selector) || selector.kind() == Kind.ENUM;
        if (!selectorComparable) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_SWITCH_SELECTOR_INVALID.name(),
                    "Switch selector type is not supported",
                    switchExpr.selector().span());
        }

        final var seenPatterns = new HashSet<String>();
        boolean hasWildcard = false;
        PbsAst.WildcardSwitchKind wildcardKind = null;
        int wildcardCount = 0;
        TypeView armType = null;

        for (final var arm : switchExpr.arms()) {
            if (arm.pattern() instanceof PbsAst.WildcardSwitchPattern wildcardPattern) {
                hasWildcard = true;
                if (wildcardPattern.kind() != PbsAst.WildcardSwitchKind.RECOVERY) {
                    wildcardCount++;
                    if (wildcardCount == 1) {
                        wildcardKind = wildcardPattern.kind();
                    } else if (wildcardKind != wildcardPattern.kind()) {
                        Diagnostics.error(diagnostics,
                                PbsSemanticsErrors.E_SEM_SWITCH_DUPLICATE_PATTERN.name(),
                                "Switch wildcard arms cannot mix 'default' and '_'",
                                arm.pattern().span());
                    } else {
                        Diagnostics.error(diagnostics,
                                PbsSemanticsErrors.E_SEM_SWITCH_DUPLICATE_PATTERN.name(),
                                "Duplicate switch wildcard pattern",
                                arm.pattern().span());
                    }
                }
            } else {
                final var patternKey = switchPatternKey(arm.pattern());
                if (!seenPatterns.add(patternKey)) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_SWITCH_DUPLICATE_PATTERN.name(),
                            "Duplicate switch pattern",
                            arm.pattern().span());
                }
            }

            final var patternType = switchPatternType(arm.pattern(), model, diagnostics);
            if (patternType != null && !typeOps.compatible(patternType, selector)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_SWITCH_PATTERN_TYPE_MISMATCH.name(),
                        "Switch pattern is not compatible with selector type",
                        arm.pattern().span());
            }

            final var currentArmType = blockAnalysisDelegate.analyze(arm.block(), context, true);
            if (armType == null) {
                armType = currentArmType;
            } else if (!typeOps.compatible(currentArmType, armType)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_SWITCH_ARM_TYPE_MISMATCH.name(),
                        "Switch arm block types must be compatible",
                        arm.span());
            }
        }

        if (context.valueContext() && !hasWildcard) {
            var exhaustive = false;
            if (selector.kind() == Kind.ENUM) {
                final var allCases = model.enums.get(selector.name());
                if (allCases != null) {
                    final var covered = new HashSet<String>();
                    for (final var arm : switchExpr.arms()) {
                        if (arm.pattern() instanceof PbsAst.EnumCaseSwitchPattern enumCasePattern
                                && enumCasePattern.path().segments().size() == 2
                                && selector.name().equals(enumCasePattern.path().segments().getFirst())) {
                            covered.add(enumCasePattern.path().segments().get(1));
                        }
                    }
                    exhaustive = covered.containsAll(allCases);
                }
            }
            if (!exhaustive) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_SWITCH_NON_EXHAUSTIVE.name(),
                        "Switch expression in value position must be exhaustive",
                        switchExpr.span());
            }
        }

        if (armType == null) {
            return TypeView.unit();
        }
        if (context.expectedType() != null && typeOps.compatible(armType, context.expectedType())) {
            return context.expectedType();
        }
        return armType;
    }

    TypeView analyzeHandleExpression(
            final PbsAst.HandleExpr handleExpr,
            final PbsFlowExpressionContext context) {
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var resultErrorName = context.resultErrorName();
        final var sourceType = analyzeValueExpression(handleExpr.value(), context, null).type();
        if (sourceType.kind() != Kind.RESULT) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_HANDLE_NON_RESULT.name(),
                    "Handle requires result expression",
                    handleExpr.value().span());
            return TypeView.unknown();
        }
        if (resultErrorName == null) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_HANDLE_ERROR_MISMATCH.name(),
                    "Handle requires enclosing callable with result return",
                    handleExpr.span());
        }

        final var sourceErrorName = sourceType.errorType() == null ? null : sourceType.errorType().name();
        final var sourceCases = sourceErrorName == null ? Set.<String>of() : model.errors.getOrDefault(sourceErrorName, Set.of());
        final var sourcePayloadType = sourceType.inner() == null ? TypeView.unknown() : sourceType.inner();
        final var matchedCases = new HashSet<String>();
        var hasWildcard = false;

        for (final var arm : handleExpr.arms()) {
            if (arm.pattern() instanceof PbsAst.WildcardHandlePattern) {
                hasWildcard = true;
            } else if (arm.pattern() instanceof PbsAst.ErrorPathHandlePattern errorPathHandlePattern) {
                final var segments = errorPathHandlePattern.path().segments();
                if (segments.size() == 2) {
                    final var errorName = segments.getFirst();
                    final var caseName = segments.get(1);
                    if (!errorName.equals(sourceErrorName) || !sourceCases.contains(caseName)) {
                        Diagnostics.error(diagnostics,
                                PbsSemanticsErrors.E_SEM_HANDLE_ERROR_MISMATCH.name(),
                                "Handle arm pattern does not match source result error type",
                                arm.pattern().span());
                    } else if (!matchedCases.add(caseName)) {
                        Diagnostics.error(diagnostics,
                                PbsSemanticsErrors.E_SEM_HANDLE_ERROR_MISMATCH.name(),
                                "Handle arm duplicates same error case pattern",
                                arm.pattern().span());
                    }
                }
            }

            if (arm.remapTarget() != null) {
                if (!matchesTargetError(arm.remapTarget(), resultErrorName, model)) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_HANDLE_ERROR_MISMATCH.name(),
                            "Handle remap target must match enclosing callable result error type",
                            arm.remapTarget().span());
                }
                continue;
            }
            analyzeHandleBlockArm(arm, context, sourcePayloadType, model, diagnostics);
        }

        if (!hasWildcard && !sourceCases.isEmpty() && !matchedCases.containsAll(sourceCases)) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_HANDLE_ERROR_MISMATCH.name(),
                    "Handle mapping is not exhaustive for source result error cases",
                    handleExpr.span());
        }

        return sourceType.inner() == null ? TypeView.unknown() : sourceType.inner();
    }

    private void analyzeHandleBlockArm(
            final PbsAst.HandleArm arm,
            final PbsFlowExpressionContext context,
            final TypeView sourcePayloadType,
            final PbsFlowSemanticSupport.Model model,
            final DiagnosticSink diagnostics) {
        final var resultErrorName = context.resultErrorName();
        if (arm.block() == null) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_HANDLE_ARM_TERMINAL_INVALID.name(),
                    "Handle block arm must terminate with 'ok(...)' or 'err(E.case)'",
                    arm.span());
            return;
        }

        final var block = arm.block();
        final var terminal = unwrapGroup(block.tailExpression());
        if (terminal instanceof PbsAst.OkExpr okExpr) {
            final var payloadBlock = new PbsAst.Block(block.statements(), okExpr.value(), block.span());
            final var actualPayloadType = blockAnalysisDelegate.analyze(payloadBlock, context, true);
            if (!typeOps.compatible(actualPayloadType, sourcePayloadType)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_RESULT_OK_PAYLOAD_MISMATCH.name(),
                        "Handle arm 'ok(...)' payload is incompatible with source result payload type",
                        okExpr.value().span());
            }
            return;
        }

        if (terminal instanceof PbsAst.ErrExpr errExpr) {
            final var statementsOnly = new PbsAst.Block(block.statements(), null, block.span());
            blockAnalysisDelegate.analyze(statementsOnly, context, false);
            if (resultErrorName != null && !matchesTargetError(errExpr.errorPath(), resultErrorName, model)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_RESULT_ERROR_LABEL_INVALID.name(),
                        "Error label in handle arm 'err(...)' does not match enclosing result error type",
                        errExpr.errorPath().span());
            }
            return;
        }

        if (block.tailExpression() == null) {
            final var statementsOnly = new PbsAst.Block(block.statements(), null, block.span());
            blockAnalysisDelegate.analyze(statementsOnly, context, false);
        } else {
            blockAnalysisDelegate.analyze(block, context, true);
        }
        Diagnostics.error(diagnostics,
                PbsSemanticsErrors.E_SEM_HANDLE_ARM_TERMINAL_INVALID.name(),
                "Handle block arm must terminate with 'ok(...)' or 'err(E.case)'",
                block.span());
    }

    private String switchPatternKey(final PbsAst.SwitchPattern pattern) {
        if (pattern instanceof PbsAst.WildcardSwitchPattern) {
            return "<wildcard>";
        }
        if (pattern instanceof PbsAst.EnumCaseSwitchPattern enumCaseSwitchPattern) {
            return "enum:" + String.join(".", enumCaseSwitchPattern.path().segments().asList());
        }
        if (pattern instanceof PbsAst.LiteralSwitchPattern literalSwitchPattern) {
            return "lit:" + literalSwitchPattern.literal().toString();
        }
        return "<unknown>";
    }

    private TypeView switchPatternType(
            final PbsAst.SwitchPattern pattern,
            final PbsFlowSemanticSupport.Model model,
            final DiagnosticSink diagnostics) {
        if (pattern instanceof PbsAst.WildcardSwitchPattern) {
            return null;
        }
        if (pattern instanceof PbsAst.EnumCaseSwitchPattern enumCaseSwitchPattern) {
            final var segments = enumCaseSwitchPattern.path().segments();
            if (segments.size() != 2) {
                return TypeView.unknown();
            }
            final var enumName = segments.getFirst();
            final var caseName = segments.get(1);
            final var enumCases = model.enums.get(enumName);
            if (enumCases == null || !enumCases.contains(caseName)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_SWITCH_PATTERN_TYPE_MISMATCH.name(),
                        "Enum case pattern '%s.%s' does not resolve".formatted(enumName, caseName),
                        enumCaseSwitchPattern.span());
                return TypeView.unknown();
            }
            return TypeView.enumType(enumName);
        }
        if (pattern instanceof PbsAst.LiteralSwitchPattern literalSwitchPattern) {
            final var literal = literalSwitchPattern.literal();
            if (literal instanceof PbsAst.IntLiteralExpr || literal instanceof PbsAst.BoundedLiteralExpr) {
                return TypeView.intType();
            }
            if (literal instanceof PbsAst.FloatLiteralExpr) {
                return TypeView.floatType();
            }
            if (literal instanceof PbsAst.BoolLiteralExpr) {
                return TypeView.bool();
            }
            if (literal instanceof PbsAst.StringLiteralExpr) {
                return TypeView.str();
            }
        }
        return TypeView.unknown();
    }

    private boolean matchesTargetError(
            final PbsAst.ErrorPath path,
            final String resultErrorName,
            final PbsFlowSemanticSupport.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);
    }

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

    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));
    }
}
