package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;

public final class PbsLifecycleSemanticsValidator {
    private static final String ATTR_INIT_ALLOWED = "InitAllowed";

    public void validate(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        final var initFunctions = new ArrayList<PbsAst.FunctionDecl>();
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.FunctionDecl functionDecl)) {
                continue;
            }
            if (functionDecl.lifecycleMarker() == PbsAst.LifecycleMarker.NONE) {
                continue;
            }
            validateLifecycleSignature(functionDecl, diagnostics);
            if (functionDecl.lifecycleMarker() == PbsAst.LifecycleMarker.INIT) {
                initFunctions.add(functionDecl);
            }
        }

        if (initFunctions.size() > 1) {
            for (int i = 1; i < initFunctions.size(); i++) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(
                        diagnostics,
                        PbsSemanticsErrors.E_SEM_DUPLICATE_FILE_INIT.name(),
                        "A source file may declare at most one [Init] function",
                        initFunctions.get(i).span());
            }
        }

        final var hostPolicy = hostPolicy(ast, supplementalTopDecls);
        for (final var initFunction : initFunctions) {
            validateInitBody(initFunction.body(), hostPolicy, diagnostics);
        }
    }

    private void validateLifecycleSignature(
            final PbsAst.FunctionDecl functionDecl,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        final var valid = functionDecl.parameters().isEmpty()
                && functionDecl.returnKind() == PbsAst.ReturnKind.EXPLICIT_UNIT
                && functionDecl.returnType() != null
                && functionDecl.returnType().kind() == PbsAst.TypeRefKind.UNIT
                && functionDecl.resultErrorType() == null;
        if (valid) {
            return;
        }
        p.studio.compiler.source.diagnostics.Diagnostics.error(
                diagnostics,
                PbsSemanticsErrors.E_SEM_INVALID_LIFECYCLE_SIGNATURE.name(),
                "Lifecycle functions marked with [%s] must have signature 'fn name() -> void'".formatted(
                        functionDecl.lifecycleMarker() == PbsAst.LifecycleMarker.INIT ? "Init" : "Frame"),
                functionDecl.span());
    }

    private HostPolicy hostPolicy(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls) {
        final var methodsByOwner = new HashMap<String, Set<String>>();
        final var initAllowedByOwner = new HashMap<String, Set<String>>();
        collectHostPolicy(ast.topDecls(), methodsByOwner, initAllowedByOwner);
        collectHostPolicy(supplementalTopDecls, methodsByOwner, initAllowedByOwner);
        return new HostPolicy(methodsByOwner, initAllowedByOwner);
    }

    private void collectHostPolicy(
            final Iterable<PbsAst.TopDecl> topDecls,
            final Map<String, Set<String>> methodsByOwner,
            final Map<String, Set<String>> initAllowedByOwner) {
        for (final var topDecl : topDecls) {
            if (!(topDecl instanceof PbsAst.HostDecl hostDecl)) {
                continue;
            }
            for (final var signature : hostDecl.signatures()) {
                methodsByOwner.computeIfAbsent(hostDecl.name(), ignored -> new HashSet<>()).add(signature.name());
                if (hasInitAllowed(signature.attributes())) {
                    initAllowedByOwner.computeIfAbsent(hostDecl.name(), ignored -> new HashSet<>()).add(signature.name());
                }
            }
        }
    }

    private boolean hasInitAllowed(final ReadOnlyList<PbsAst.Attribute> attributes) {
        for (final var attribute : attributes) {
            if (ATTR_INIT_ALLOWED.equals(attribute.name())) {
                return true;
            }
        }
        return false;
    }

    private void validateInitBody(
            final PbsAst.Block block,
            final HostPolicy hostPolicy,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        if (block == null) {
            return;
        }
        for (final var statement : block.statements()) {
            validateStatement(statement, hostPolicy, diagnostics);
        }
        validateExpression(block.tailExpression(), hostPolicy, diagnostics);
    }

    private void validateStatement(
            final PbsAst.Statement statement,
            final HostPolicy hostPolicy,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        if (statement instanceof PbsAst.LetStatement letStatement) {
            validateExpression(letStatement.initializer(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.AssignStatement assignStatement) {
            validateExpression(assignStatement.value(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.ReturnStatement returnStatement) {
            validateExpression(returnStatement.value(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.IfStatement ifStatement) {
            validateExpression(ifStatement.condition(), hostPolicy, diagnostics);
            validateInitBody(ifStatement.thenBlock(), hostPolicy, diagnostics);
            validateStatement(ifStatement.elseIf(), hostPolicy, diagnostics);
            validateInitBody(ifStatement.elseBlock(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.ForStatement forStatement) {
            validateExpression(forStatement.fromExpression(), hostPolicy, diagnostics);
            validateExpression(forStatement.untilExpression(), hostPolicy, diagnostics);
            validateExpression(forStatement.stepExpression(), hostPolicy, diagnostics);
            validateInitBody(forStatement.body(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.WhileStatement whileStatement) {
            validateExpression(whileStatement.condition(), hostPolicy, diagnostics);
            validateInitBody(whileStatement.body(), hostPolicy, diagnostics);
            return;
        }
        if (statement instanceof PbsAst.ExpressionStatement expressionStatement) {
            validateExpression(expressionStatement.expression(), hostPolicy, diagnostics);
        }
    }

    private void validateExpression(
            final PbsAst.Expression expression,
            final HostPolicy hostPolicy,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        if (expression == null) {
            return;
        }
        if (expression instanceof PbsAst.CallExpr callExpr) {
            validateHostCall(callExpr, hostPolicy, diagnostics);
            validateExpression(callExpr.callee(), hostPolicy, diagnostics);
            for (final var argument : callExpr.arguments()) {
                validateExpression(argument, hostPolicy, diagnostics);
            }
            return;
        }
        if (expression instanceof PbsAst.ApplyExpr applyExpr) {
            validateExpression(applyExpr.callee(), hostPolicy, diagnostics);
            validateExpression(applyExpr.argument(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.MemberExpr memberExpr) {
            validateExpression(memberExpr.receiver(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.UnaryExpr unaryExpr) {
            validateExpression(unaryExpr.expression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.BinaryExpr binaryExpr) {
            validateExpression(binaryExpr.left(), hostPolicy, diagnostics);
            validateExpression(binaryExpr.right(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.GroupExpr groupExpr) {
            validateExpression(groupExpr.expression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.NewExpr newExpr) {
            for (final var argument : newExpr.arguments()) {
                validateExpression(argument, hostPolicy, diagnostics);
            }
            return;
        }
        if (expression instanceof PbsAst.IfExpr ifExpr) {
            validateExpression(ifExpr.condition(), hostPolicy, diagnostics);
            validateInitBody(ifExpr.thenBlock(), hostPolicy, diagnostics);
            validateExpression(ifExpr.elseExpression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.SwitchExpr switchExpr) {
            validateExpression(switchExpr.selector(), hostPolicy, diagnostics);
            for (final var arm : switchExpr.arms()) {
                validateInitBody(arm.block(), hostPolicy, diagnostics);
            }
            return;
        }
        if (expression instanceof PbsAst.HandleExpr handleExpr) {
            validateExpression(handleExpr.value(), hostPolicy, diagnostics);
            for (final var arm : handleExpr.arms()) {
                validateInitBody(arm.block(), hostPolicy, diagnostics);
            }
            return;
        }
        if (expression instanceof PbsAst.ElseExpr elseExpr) {
            validateExpression(elseExpr.optionalExpression(), hostPolicy, diagnostics);
            validateExpression(elseExpr.fallbackExpression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.PropagateExpr propagateExpr) {
            validateExpression(propagateExpr.expression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.BindExpr bindExpr) {
            validateExpression(bindExpr.contextExpression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.AsExpr asExpr) {
            validateExpression(asExpr.expression(), hostPolicy, diagnostics);
            return;
        }
        if (expression instanceof PbsAst.BlockExpr blockExpr) {
            validateInitBody(blockExpr.block(), hostPolicy, diagnostics);
        }
    }

    private void validateHostCall(
            final PbsAst.CallExpr callExpr,
            final HostPolicy hostPolicy,
            final p.studio.compiler.source.diagnostics.DiagnosticSink diagnostics) {
        if (!(callExpr.callee() instanceof PbsAst.MemberExpr memberExpr)) {
            return;
        }
        if (!(memberExpr.receiver() instanceof PbsAst.IdentifierExpr receiverIdentifier)) {
            return;
        }
        if (!hostPolicy.isHostMethod(receiverIdentifier.name(), memberExpr.memberName())) {
            return;
        }
        if (hostPolicy.isInitAllowed(receiverIdentifier.name(), memberExpr.memberName())) {
            return;
        }
        p.studio.compiler.source.diagnostics.Diagnostics.error(
                diagnostics,
                PbsSemanticsErrors.E_SEM_INIT_HOST_CALL_NOT_ALLOWED.name(),
                "Host call '%s.%s' is not allowed during [Init] without InitAllowed".formatted(
                        receiverIdentifier.name(),
                        memberExpr.memberName()),
                callExpr.span());
    }

    private record HostPolicy(
            Map<String, Set<String>> methodsByOwner,
            Map<String, Set<String>> initAllowedByOwner) {
        boolean isHostMethod(final String owner, final String method) {
            return methodsByOwner.getOrDefault(owner, Set.of()).contains(method);
        }

        boolean isInitAllowed(final String owner, final String method) {
            return initAllowedByOwner.getOrDefault(owner, Set.of()).contains(method);
        }
    }
}
