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

final class PbsFlowAssignmentAnalyzer {
    @FunctionalInterface
    interface BlockAnalysisDelegate {
        TypeView analyze(
                PbsAst.Block block,
                Scope outerScope,
                TypeView returnType,
                String resultErrorName,
                TypeView receiverType,
                Model model,
                p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics,
                boolean valueContext);
    }

    private record AssignmentTargetResolution(
            TypeView type,
            boolean assignable) {
    }

    private final PbsFlowTypeOps typeOps;
    private final PbsFlowExpressionAnalyzer expressionAnalyzer;
    private final BlockAnalysisDelegate blockAnalysisDelegate;

    PbsFlowAssignmentAnalyzer(
            final PbsFlowTypeOps typeOps,
            final PbsFlowExpressionAnalyzer expressionAnalyzer,
            final BlockAnalysisDelegate blockAnalysisDelegate) {
        this.typeOps = typeOps;
        this.expressionAnalyzer = expressionAnalyzer;
        this.blockAnalysisDelegate = blockAnalysisDelegate;
    }

    void analyzeAssignmentStatement(
            final PbsAst.AssignStatement assignStatement,
            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 target = resolveAssignmentTarget(assignStatement.target(), context);
        final var expectedType = target.type();
        final var valueType = expressionAnalyzer.analyzeExpression(
                assignStatement.value(),
                scope,
                expectedType,
                returnType,
                resultErrorName,
                receiverType,
                model,
                diagnostics,
                ExprUse.VALUE,
                true,
                blockAnalysisDelegate::analyze).type();

        if (!target.assignable()) {
            return;
        }

        if (!assignmentCompatible(assignStatement.operator(), target.type(), valueType)) {
            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_ASSIGN_TYPE_MISMATCH.name(),
                    "Assigned value type is not compatible with assignment target",
                    assignStatement.span());
        }
    }

    private AssignmentTargetResolution resolveAssignmentTarget(
            final PbsAst.LValue lValue,
            final PbsFlowBodyContext context) {
        final var scope = context.scope();
        final var receiverType = context.receiverType();
        final var model = context.model();
        final var diagnostics = context.diagnostics();
        final var rootName = lValue.rootName();
        final var rootIsThis = "this".equals(rootName);
        if (lValue.pathSegments().isEmpty()) {
            if (rootIsThis) {
                if (receiverType == null) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_INVALID_THIS_CONTEXT.name(),
                            "Invalid 'this' usage outside struct/service methods and constructors",
                            lValue.span());
                }
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_ASSIGN_TARGET_NOT_ASSIGNABLE.name(),
                        "Assignment target is not assignable",
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }

            final var localType = scope.resolve(rootName);
            if (localType != null) {
                if (!scope.isMutable(rootName)) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_ASSIGN_TARGET_NOT_ASSIGNABLE.name(),
                            "Cannot assign to immutable local '%s'".formatted(rootName),
                            lValue.span());
                    return new AssignmentTargetResolution(localType, false);
                }
                return new AssignmentTargetResolution(localType, true);
            }

            final var globalType = model.globalTypes.get(rootName);
            if (globalType != null) {
                return new AssignmentTargetResolution(globalType, true);
            }

            if (isKnownNonAssignableRoot(rootName, model)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_ASSIGN_TARGET_NOT_ASSIGNABLE.name(),
                        "Assignment target is not assignable",
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }

            Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_ASSIGN_TARGET_UNRESOLVED.name(),
                    "Assignment target '%s' does not resolve".formatted(rootName),
                    lValue.span());
            return new AssignmentTargetResolution(TypeView.unknown(), false);
        }

        TypeView currentType;
        var currentReceiverIsThis = rootIsThis;
        if (rootIsThis) {
            if (receiverType == null) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_THIS_CONTEXT.name(),
                        "Invalid 'this' usage outside struct/service methods and constructors",
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }
            currentType = receiverType;
        } else {
            final var localType = scope.resolve(rootName);
            if (localType != null) {
                currentType = localType;
            } else if (model.globalTypes.containsKey(rootName)) {
                currentType = model.globalTypes.get(rootName);
            } else if (model.serviceSingletons.containsKey(rootName)) {
                currentType = model.serviceSingletons.get(rootName);
            } else if (model.constTypes.containsKey(rootName)) {
                currentType = model.constTypes.get(rootName);
            } else {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_ASSIGN_TARGET_UNRESOLVED.name(),
                        "Assignment target '%s' does not resolve".formatted(rootName),
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }
        }

        for (int i = 0; i < lValue.pathSegments().size(); i++) {
            final var segment = lValue.pathSegments().get(i);
            final var isLastSegment = i == lValue.pathSegments().size() - 1;
            if (currentType.kind() != Kind.STRUCT) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                        "Assignment target segment '%s' is not a struct field".formatted(segment),
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }

            final var struct = model.structs.get(currentType.name());
            if (struct == null) {
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }
            final var field = struct.fields().get(segment);
            if (field == null) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                        "Struct member '%s' does not exist".formatted(segment),
                        lValue.span());
                return new AssignmentTargetResolution(TypeView.unknown(), false);
            }

            final var ownerContext = receiverType != null
                    && receiverType.kind() == Kind.STRUCT
                    && receiverType.name().equals(currentType.name());
            if (isLastSegment) {
                if (!canWriteStructField(field, ownerContext, currentReceiverIsThis)) {
                    Diagnostics.error(diagnostics,
                            PbsSemanticsErrors.E_SEM_FIELD_WRITE_ACCESS_DENIED.name(),
                            "Write access to struct field '%s' is not permitted".formatted(segment),
                            lValue.span());
                    return new AssignmentTargetResolution(field.type(), false);
                }
                return new AssignmentTargetResolution(field.type(), true);
            }

            if (!canReadStructField(field, ownerContext, currentReceiverIsThis)) {
                Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_FIELD_READ_ACCESS_DENIED.name(),
                        "Read access to struct field '%s' is not permitted".formatted(segment),
                        lValue.span());
                return new AssignmentTargetResolution(field.type(), false);
            }

            currentType = field.type();
            currentReceiverIsThis = false;
        }

        return new AssignmentTargetResolution(TypeView.unknown(), false);
    }

    private boolean assignmentCompatible(
            final PbsAst.AssignOperator operator,
            final TypeView targetType,
            final TypeView valueType) {
        if (operator == PbsAst.AssignOperator.ASSIGN) {
            return typeOps.compatible(valueType, targetType);
        }

        final String binaryOperator = switch (operator) {
            case ADD_ASSIGN -> "+";
            case SUB_ASSIGN -> "-";
            case MUL_ASSIGN -> "*";
            case DIV_ASSIGN -> "/";
            case MOD_ASSIGN -> "%";
            default -> throw new IllegalStateException("Unexpected value: " + operator);
        };
        final var resultType = typeOps.inferBinaryResult(binaryOperator, targetType, valueType);
        return typeOps.compatible(resultType, targetType);
    }

    private boolean canReadStructField(
            final PbsFlowSemanticSupport.StructFieldInfo field,
            final boolean ownerContext,
            final boolean receiverIsThis) {
        if (field.isPublic()) {
            return true;
        }
        return ownerContext && receiverIsThis;
    }

    private boolean canWriteStructField(
            final PbsFlowSemanticSupport.StructFieldInfo field,
            final boolean ownerContext,
            final boolean receiverIsThis) {
        if (!field.isPublic()) {
            return ownerContext && receiverIsThis;
        }
        if (ownerContext) {
            return true;
        }
        return field.isMutable();
    }

    private boolean isKnownNonAssignableRoot(
            final String rootName,
            final Model model) {
        return model.serviceSingletons.containsKey(rootName)
                || model.constTypes.containsKey(rootName)
                || model.topLevelCallables.containsKey(rootName)
                || model.callbacks.containsKey(rootName)
                || model.enums.containsKey(rootName)
                || model.errors.containsKey(rootName)
                || model.structs.containsKey(rootName)
                || model.services.containsKey(rootName)
                || model.contracts.containsKey(rootName);
    }
}
